perf(observations): reduce batch pipeline copies (#1319)

src/unilab/managers/_noise/noise_cfg.py

@@ -68,14 +68,21 @@ def apply(self, data: np.ndarray, *, rng: np.random.Generator | None = None) ->
         n_min = self._as_array(self.n_min, data.dtype)
         n_max = self._as_array(self.n_max, data.dtype)
 
-        # Generate uniform noise in [0, 1) and scale to [n_min, n_max).
+        # Generate uniform noise in [0, 1) and transform the generated array
+        # in place.  The pre-existing expression allocated one array for each
+        # multiply, add, and final data operation; keeping the same float32
+        # cast and ufunc order preserves bit-level results while returning the
+        # transformed noise buffer itself.
         noise = rng.random(data.shape).astype(data.dtype, copy=False)
-        noise = noise * (n_max - n_min) + n_min
+        np.multiply(noise, n_max - n_min, out=noise)
+        np.add(noise, n_min, out=noise)
 
         if self.operation == "add":
-            return data + noise
+            np.add(data, noise, out=noise)
+            return noise
         elif self.operation == "scale":
-            return data * noise
+            np.multiply(data, noise, out=noise)
+            return noise
         elif self.operation == "abs":
             return noise
         else:

src/unilab/managers/observation_manager.py

@@ -335,7 +335,7 @@ def _row_env_ids(mask: np.ndarray) -> list[int]:
             )
 
         # Sanitize (applies to both "warn" and "sanitize" policies).
-        return np.nan_to_num(tensor, nan=0.0, posinf=0.0, neginf=0.0)
+        return np.nan_to_num(tensor, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
 
     def compute(
         self,
@@ -381,6 +381,17 @@ def compute_group(
         group_term_names = self._group_obs_term_names[group_name]
         group_obs: dict[str, np.ndarray] = {}
         obs_terms = zip(group_term_names, self._group_obs_term_cfgs[group_name], strict=False)
+        # In the strict default policy a finite result is by far the common
+        # case.  For concatenated groups, scan the assembled output once and
+        # only inspect individual slices when an error is actually found; this
+        # retains the per-term diagnostic while removing repeated full scans
+        # from the hot path.
+        defer_error_nan_check = (
+            self._group_obs_concatenate[group_name]
+            and not self._group_obs_temporal[group_name]
+            and group_cfg.nan_check_per_term
+            and group_cfg.nan_policy == "error"
+        )
         # Reset path (issue #1259 R2): when no term in this group uses delay or
         # history buffers, everything downstream of the term call is row
         # independent, so only the reset rows are processed. Term calls and
@@ -409,10 +420,28 @@ def compute_group(
                 # NoiseModel.__call__ likewise returns a new array.
                 obs = self._group_obs_class_instances[group_name][term_name](obs)
                 fresh = True
-            if not row_scoped and not fresh:
-                # Terms may return backend/command-owned buffers; copy before the
-                # in-place clip/scale below. Skipped when noise already produced
-                # a fresh array (issue #1296).
+            sanitizes_per_term = group_cfg.nan_check_per_term and group_cfg.nan_policy in (
+                "warn",
+                "sanitize",
+            )
+            exposes_term_output = (
+                not self._group_obs_concatenate[group_name]
+                and term_cfg.delay_max_lag == 0
+                and term_cfg.history_length == 0
+            )
+            if (
+                not row_scoped
+                and not fresh
+                and (
+                    term_cfg.clip is not None
+                    or term_cfg.scale is not None
+                    or sanitizes_per_term
+                    or exposes_term_output
+                )
+            ):
+                # Concatenation and temporal buffers already copy their inputs.
+                # Only take a defensive copy when this pipeline may mutate the
+                # term or expose it directly to callers.
                 obs = obs.copy()
             if row_scoped:
                 # Fresh row copy; safe for the in-place clip/scale below.
@@ -425,7 +454,11 @@ def compute_group(
                 np.multiply(obs, scale, out=obs)
 
             # Check for NaN/Inf before delay/history buffers (per-term checking).
-            if group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled":
+            if (
+                group_cfg.nan_check_per_term
+                and group_cfg.nan_policy != "disabled"
+                and not defer_error_nan_check
+            ):
                 obs = self._check_and_handle_nans(
                     obs,
                     context=f"{group_name}/{term_name}",
@@ -474,6 +507,33 @@ def compute_group(
             result = np.concatenate(
                 list(group_obs.values()), axis=self._group_obs_concatenate_dim[group_name]
             )
+            if defer_error_nan_check:
+                finite = np.isfinite(result)
+                if not finite.all():
+                    axis = self._group_obs_concatenate_dim[group_name]
+                    axis = axis if axis >= 0 else result.ndim + axis
+                    offset = 0
+                    for term_name, term_dims in zip(
+                        group_term_names,
+                        self._group_obs_term_dim[group_name],
+                        strict=True,
+                    ):
+                        width = int(term_dims[axis - 1])
+                        selectors = [slice(None)] * result.ndim
+                        selectors[axis] = slice(offset, offset + width)
+                        selector = tuple(selectors)
+                        if not finite[selector].all():
+                            # Reuse the established diagnostic path so the
+                            # error still names the first offending term and
+                            # reports reset-row IDs when applicable.
+                            self._check_and_handle_nans(
+                                result[selector],
+                                context=f"{group_name}/{term_name}",
+                                policy=group_cfg.nan_policy,
+                                env_ids=env_ids if row_scoped else None,
+                            )
+                            break
+                        offset += width
             # Final check for concatenated result (non-per-term checking).
             if not group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled":
                 result = self._check_and_handle_nans(

tests/managers/test_observation_buffers_noise.py

@@ -95,6 +95,28 @@ def test_noise_configs_use_supplied_generator() -> None:
     np.testing.assert_array_equal(ConstantNoiseCfg(bias=2.0, operation="abs").apply(data), 2.0)
 
 
+@pytest.mark.parametrize("operation", ["add", "scale", "abs"])
+def test_uniform_noise_inplace_matches_reference_expression(operation: str) -> None:
+    data = np.arange(24, dtype=np.float32).reshape(8, 3)
+    n_min = np.asarray([-0.2, -0.1, -0.05], dtype=np.float32)
+    n_max = np.asarray([0.3, 0.4, 0.5], dtype=np.float32)
+    cfg = UniformNoiseCfg(n_min=tuple(n_min), n_max=tuple(n_max), operation=operation)
+
+    reference_rng = np.random.default_rng(1702)
+    unit = reference_rng.random(data.shape).astype(data.dtype, copy=False)
+    noise = unit * (n_max - n_min) + n_min
+    if operation == "add":
+        expected = data + noise
+    elif operation == "scale":
+        expected = data * noise
+    else:
+        expected = noise
+
+    actual = cfg.apply(data, rng=np.random.default_rng(1702))
+    np.testing.assert_array_equal(actual, expected)
+    np.testing.assert_array_equal(data, np.arange(24, dtype=np.float32).reshape(8, 3))
+
+
 def test_additive_bias_noise_supports_scalar_terms() -> None:
     from unilab.managers._noise import NoiseModelWithAdditiveBias
 
@@ -143,6 +165,86 @@ def test_observation_groups_pipeline_order_and_history(fake_env: FakeEnv) -> Non
     assert manager.get_active_iterable_terms(0)[0][0] == "policy-state"
 
 
+def test_concatenated_result_owns_each_result_and_protects_term_buffers(
+    fake_env: FakeEnv,
+) -> None:
+    source = fake_env.obs
+    manager = ObservationManager(
+        {
+            "policy": ObservationGroupCfg(
+                terms={
+                    "source": ObservationTermCfg(func=lambda env: env.obs),
+                    "constant": ObservationTermCfg(
+                        func=lambda env: np.ones((env.num_envs, 1), dtype=np.float32)
+                    ),
+                }
+            )
+        },
+        fake_env,
+    )
+
+    first = manager.compute(update_history=True)["policy"]
+    assert isinstance(first, np.ndarray)
+    first_address = first.ctypes.data
+    expected_first = np.concatenate((source.copy(), np.ones((fake_env.num_envs, 1))), axis=1)
+    np.testing.assert_array_equal(first, expected_first)
+
+    fake_env.obs += 100.0
+    second = manager.compute(update_history=True)["policy"]
+    assert isinstance(second, np.ndarray)
+    assert second.ctypes.data != first_address
+    np.testing.assert_array_equal(first, expected_first)
+    np.testing.assert_array_equal(second[:, :2], fake_env.obs)
+    assert not np.shares_memory(second, fake_env.obs)
+
+
+def test_concatenated_nan_sanitize_does_not_mutate_term_owned_input(
+    fake_env: FakeEnv,
+) -> None:
+    source = fake_env.obs.copy()
+    source[0, 0] = np.nan
+    manager = ObservationManager(
+        {
+            "policy": ObservationGroupCfg(
+                terms={"state": ObservationTermCfg(func=lambda env: source)},
+                nan_policy="sanitize",
+                nan_check_per_term=False,
+            )
+        },
+        fake_env,
+    )
+
+    result = manager.compute(update_history=True)["policy"]
+    assert isinstance(result, np.ndarray)
+    assert np.isfinite(result).all()
+    assert np.isnan(source[0, 0])
+
+
+def test_concatenated_nan_error_still_identifies_offending_term(fake_env: FakeEnv) -> None:
+    def invalid(env: FakeEnv) -> np.ndarray:
+        result = np.ones((env.num_envs, 1), dtype=np.float32)
+        result[2, 0] = np.nan
+        return result
+
+    manager = ObservationManager(
+        {
+            "policy": ObservationGroupCfg(
+                terms={
+                    "finite": ObservationTermCfg(func=lambda env: env.obs),
+                    "invalid": ObservationTermCfg(func=invalid),
+                }
+            )
+        },
+        fake_env,
+    )
+
+    with pytest.raises(
+        ValueError,
+        match=r"NaN detected.*'policy/invalid'.*environments: \[2\]",
+    ):
+        manager.compute(update_history=True)
+
+
 def test_observation_noise_model_delay_and_seed_reproducibility() -> None:
     cfg = {
         "policy": ObservationGroupCfg(