fix(fastsac): preserve final observations and stabilize actor

motrix_env_core/src/motrix_env_core/array/env.py

@@ -33,6 +33,18 @@ def value_or_policy(self) -> np.ndarray:
         return self.policy if self.value is None else self.value
 
 
+@dataclass
+class FinalObs:
+    """Owned pre-reset observations of the current transition's truncated rows.
+
+    Rows in ``obs`` align with ``env_ids``. Arrays remain valid across steps;
+    consumers must treat them as read-only. ``obs.value`` stays optional.
+    """
+
+    env_ids: np.ndarray
+    obs: NpObs
+
+
 @dataclass
 class ArrayEnvState:
     obs: NpObs
@@ -50,6 +62,8 @@ class ArrayEnvState:
     # reflect the latest step: retain values across steps through
     # :meth:`process_metrics`, which returns a fresh dict of reduced scalars.
     metrics: dict[str, float | np.ndarray] = dataclasses.field(default_factory=dict)
+    # None when no row truncated or the implementation does not capture finals.
+    final_obs: FinalObs | None = None
 
     @property
     def done(self) -> np.ndarray:
@@ -234,6 +248,15 @@ def reset(self, env_ids: np.ndarray) -> None:
         writes have completed.
         """
 
+    def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs | None:
+        """Capture truncated rows before reset; default implementation is unsupported.
+
+        Called after transition and truncation updates, before any reset write.
+        Implementations must leave normal observations and task state unchanged
+        and return owned row arrays, without advancing observation history.
+        """
+        return None
+
     def _max_episode_steps(self) -> int | None:
         return self._cfg.max_episode_steps
 
@@ -279,6 +302,8 @@ def physics_step(self):
         """Advance the simulator by the configured substeps."""
 
     def _prev_physics_step(self):
+        if self._state.final_obs is not None:
+            self._state = self._state.replace(final_obs=None)
         state = self._state
         state.reward.fill(0.0)
         state.terminated.fill(False)
@@ -302,6 +327,7 @@ def step(self, actions: np.ndarray) -> ArrayEnvState:
 
         Timing: ``apply_action`` -> ``physics_step`` -> ``compute_transition``
         (the only full simulator-data refresh) -> episode bookkeeping ->
+        optional pre-reset ``final_obs`` capture for truncated rows ->
         ``reset(env_ids)`` for done rows (may refresh selected rows) ->
         ``compute_observation`` (whole batch, pure construction).
         """
@@ -319,6 +345,10 @@ def step(self, actions: np.ndarray) -> ArrayEnvState:
             self._state = self.compute_transition(self._state)
         self._state.episode_steps += 1
         self._update_truncate()
+        trunc_ids = np.flatnonzero(self._state.truncated)
+        if trunc_ids.size:
+            with self.perf.scope("final_obs"):
+                self._state = self._state.replace(final_obs=self._capture_final_obs(trunc_ids))
         with self.perf.scope("reset"):
             self._reset_done_envs()
         with self.perf.scope("observation"):
@@ -327,4 +357,4 @@ def step(self, actions: np.ndarray) -> ArrayEnvState:
         return self._state
 
 
-__all__ = ["ArrayEnv", "ArrayEnvState", "EnvCfgType", "NpObs"]
+__all__ = ["ArrayEnv", "ArrayEnvState", "EnvCfgType", "FinalObs", "NpObs"]

motrix_env_core/src/motrix_env_core/numba/manager/env.py

@@ -13,7 +13,7 @@
 import gymnasium as gym
 import numpy as np
 
-from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, EnvCfgType, NpObs
+from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, EnvCfgType, FinalObs, NpObs
 from motrix_env_core.base import EnvCfg, ObsSpace
 from motrix_env_core.config import configclass
 from motrix_env_core.config.sim_reset import ManagerResetCfg, ResetTermCfg
@@ -285,6 +285,21 @@ class _ObsBufferSlot:
     observation: NpObs
     kernel_outputs: NumbaKernelOutputs
 
+    @classmethod
+    def bind(cls, observation: NpObs, reward: np.ndarray, terminated: np.ndarray) -> _ObsBufferSlot:
+        value = observation.value
+        if value is None:
+            value = np.empty((observation.policy.shape[0], 0), dtype=observation.policy.dtype)
+        return cls(
+            observation=observation,
+            kernel_outputs=NumbaKernelOutputs(
+                policy_obs=observation.policy,
+                value_obs=value,
+                reward=reward,
+                terminated=terminated,
+            ),
+        )
+
 
 @dataclass
 class SwapObsBuffer:
@@ -311,21 +326,7 @@ def create(
         reward: np.ndarray,
         terminated: np.ndarray,
     ) -> SwapObsBuffer:
-        def make_slot(observation: NpObs) -> _ObsBufferSlot:
-            value = observation.value
-            if value is None:
-                value = np.empty((observation.policy.shape[0], 0), dtype=observation.policy.dtype)
-            return _ObsBufferSlot(
-                observation=observation,
-                kernel_outputs=NumbaKernelOutputs(
-                    policy_obs=observation.policy,
-                    value_obs=value,
-                    reward=reward,
-                    terminated=terminated,
-                ),
-            )
-
-        return cls((make_slot(first), make_slot(second)))
+        return cls((_ObsBufferSlot.bind(first, reward, terminated), _ObsBufferSlot.bind(second, reward, terminated)))
 
     @property
     def current(self) -> _ObsBufferSlot:
@@ -370,6 +371,7 @@ def __init__(self, cfg: EnvCfgType, num_envs: int = 1, backend: str | None = Non
         self.metrics: dict[str, np.ndarray] = {}
         self._kernel_outputs: NumbaKernelOutputs | None = None
         self._observation_buffer: SwapObsBuffer | None = None
+        self._final_obs_slot: _ObsBufferSlot | None = None
         # Persistent kernel-bound sim-reset-request buffer
         # (``ctx.sim_reset_requested``), cleared before each physics step and
         # consumed by _reset_done_envs.
@@ -509,6 +511,8 @@ def init_state(self) -> ArrayEnvState:
             state.terminated,
         )
         self._kernel_outputs = self._observation_buffer.current.kernel_outputs
+        final_obs = self._zeros_obs(self.observation_space)
+        self._final_obs_slot = _ObsBufferSlot.bind(final_obs, state.reward, state.terminated)
         self._refresh_sim_reads()
         inputs = self._kernel_inputs
         self._validate_kernel_context(inputs)
@@ -548,6 +552,25 @@ def _execute_observe_kernel(self, inputs: tuple[Any, ...]) -> None:
         assert self._kernel_outputs is not None
         self._task_program.observe_kernel(inputs, self._kernel_outputs)
 
+    def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs:
+        """Observe pre-reset inputs after command on_transition.
+
+        The compiled observe pass is batch-wide. Dedicated outputs preserve the
+        normal swap buffer; advanced indexing creates owned compact snapshots.
+        """
+        assert self._task_program is not None
+        assert self._final_obs_slot is not None
+        slot = self._final_obs_slot
+        self._task_program.observe_kernel(self._kernel_inputs, slot.kernel_outputs)
+        obs = slot.observation
+        return FinalObs(
+            env_ids=env_ids,
+            obs=NpObs(
+                policy=obs.policy[env_ids],
+                value=None if obs.value is None else obs.value[env_ids],
+            ),
+        )
+
     def _execute_task_kernel(self, inputs: tuple[Any, ...]) -> None:
         """Evaluate and observe without advancing host-side commands."""
         self._execute_evaluate_kernel(inputs)

motrix_env_core/src/motrix_env_core/numba/manager/observations.py

@@ -30,7 +30,12 @@ def __call__(self, ctx: BuildContext) -> ObsTerm:
 
 @dataclass(frozen=True, slots=True, init=False)
 class ObsTerm(BaseTerm):
-    """Observation dispatch term with a fixed output width."""
+    """Pure observation dispatch with a fixed output width.
+
+    Read context and write only the supplied output, except observation noise
+    may consume ctx.rand. Do not advance history or commands: truncated steps
+    may observe again before reset, with the RNG restored after that snapshot.
+    """
 
     size: int
 

motrix_env_core/tests/test_array_env.py

@@ -1,12 +1,90 @@
 # Copyright Motphys Technology Co., Ltd. 2025, 2026
 # SPDX-License-Identifier: Apache-2.0
 
-"""Regression tests for the ArrayEnv frontend hierarchy (issue #222 MR1)."""
+"""Regression tests for the ArrayEnv hierarchy and final-observation lifecycle."""
 
-from motrix_env_core.array.env import ArrayEnv
+import gymnasium as gym
+import numpy as np
+import pytest
+
+from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, FinalObs, NpObs
+from motrix_env_core.base import EnvCfg, ObsSpace
+from motrix_env_core.config.scene import SceneCfg
 from motrix_env_core.direct.env import DirectEnv
 from motrix_env_core.numba.manager.env import ManagerEnv
 from motrix_env_core.registry import _infer_sim_backend
+from motrix_env_core.sim.backend import RenderConfig, SimRenderer
+
+
+class _CounterEnv(ArrayEnv):
+    """Exercise the frontend lifecycle without constructing a simulator backend."""
+
+    def __init__(self, *, with_value: bool = False):
+        super().__init__(EnvCfg(scene=SceneCfg(), ctrl_dt=1.0, max_episode_seconds=3.0), num_envs=4)
+        policy_space = gym.spaces.Box(-np.inf, np.inf, shape=(1,), dtype=np.float32)
+        value_space = gym.spaces.Box(-np.inf, np.inf, shape=(2,), dtype=np.float32) if with_value else None
+        self._observation_space = ObsSpace(policy=policy_space, value=value_space)
+        self._action_space = gym.spaces.Box(-1.0, 1.0, shape=(1,), dtype=np.float32)
+        self.counters = np.zeros(self.num_envs, dtype=np.float32)
+        self.next_terminated = np.zeros(self.num_envs, dtype=bool)
+        self.reset_rows: list[np.ndarray] = []
+        self.physics_steps = 0
+        self.observation_calls = 0
+
+    @property
+    def observation_space(self) -> ObsSpace:
+        return self._observation_space
+
+    @property
+    def action_space(self) -> gym.spaces.Box:
+        return self._action_space
+
+    def reset(self, env_ids: np.ndarray) -> None:
+        self.reset_rows.append(env_ids.copy())
+        self.counters[env_ids] = 0.0
+
+    def apply_action(self, actions: np.ndarray, state: ArrayEnvState) -> ArrayEnvState:
+        return state
+
+    def physics_step(self) -> None:
+        self.physics_steps += 1
+        self.counters += 1.0
+
+    def compute_transition(self, state: ArrayEnvState) -> ArrayEnvState:
+        state.reward[:] = self.counters
+        state.terminated[:] = self.next_terminated
+        return state
+
+    def compute_observation(self, state: ArrayEnvState) -> ArrayEnvState:
+        self.observation_calls += 1
+        state.obs.policy[:, 0] = self.counters
+        if state.obs.value is not None:
+            state.obs.value[:, 0] = self.counters
+            state.obs.value[:, 1] = -self.counters
+        return state
+
+    def create_renderer(self, config: RenderConfig) -> SimRenderer:
+        raise NotImplementedError("The counter environment has no simulator or renderer")
+
+
+class _FinalCounterEnv(_CounterEnv):
+    def __init__(self, *, with_value: bool = False):
+        super().__init__(with_value=with_value)
+        self.capture_rows: list[np.ndarray] = []
+        self.capture_episode_steps: list[np.ndarray] = []
+        self.capture_rewards: list[np.ndarray] = []
+
+    def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs:
+        self.capture_rows.append(env_ids.copy())
+        self.capture_episode_steps.append(self.state.episode_steps[env_ids].copy())
+        self.capture_rewards.append(self.state.reward[env_ids].copy())
+        policy = self.counters[env_ids, None].copy()
+        value = np.concatenate((policy, -policy), axis=1) if self.has_value_observation else None
+        return FinalObs(env_ids=env_ids.copy(), obs=NpObs(policy=policy, value=value))
+
+
+def _step(env: _CounterEnv) -> ArrayEnvState:
+    return env.step(np.zeros((env.num_envs, 1), dtype=np.float32))
 
 
 def test_np_env_inherits_array_env() -> None:
@@ -21,3 +99,115 @@ def test_manager_env_is_array_env_sibling_of_np_env() -> None:
 def test_registry_infers_np_backend_for_manager_env() -> None:
     assert _infer_sim_backend(ManagerEnv) == "np"
     assert _infer_sim_backend(DirectEnv) == "np"
+
+
+@pytest.mark.parametrize("with_value", [False, True])
+def test_final_obs_absent_on_init_and_nontruncated_step(with_value: bool) -> None:
+    env = _FinalCounterEnv(with_value=with_value)
+    initial = env.init_state()
+
+    assert initial.final_obs is None
+    np.testing.assert_array_equal(initial.obs.policy, np.zeros((4, 1)))
+    np.testing.assert_array_equal(env.reset_rows[0], np.arange(4))
+    assert env.physics_steps == 0
+    assert env.observation_calls == 1
+    assert env.capture_rows == []
+
+    state = _step(env)
+
+    assert state.final_obs is None
+    assert not state.truncated.any()
+    assert not state.terminated.any()
+    np.testing.assert_array_equal(state.episode_steps, np.ones(4))
+    np.testing.assert_array_equal(state.obs.policy, np.ones((4, 1)))
+    assert env.physics_steps == 1
+    assert env.observation_calls == 2
+    assert len(env.reset_rows) == 1
+    assert env.capture_rows == []
+
+
+@pytest.mark.parametrize("with_value", [False, True])
+def test_final_obs_contains_only_truncated_rows_before_auto_reset(with_value: bool) -> None:
+    env = _FinalCounterEnv(with_value=with_value)
+    env.init_state()
+    env.state.episode_steps[:] = [2, 0, 2, 0]
+    env.counters[:] = [10, 20, 30, 40]
+    env.next_terminated[1] = True
+
+    state = _step(env)
+
+    final = state.final_obs
+    assert isinstance(final, FinalObs)
+    assert isinstance(final.obs, NpObs)
+    np.testing.assert_array_equal(state.truncated, [True, False, True, False])
+    np.testing.assert_array_equal(state.terminated, [False, True, False, False])
+    np.testing.assert_array_equal(final.env_ids, [0, 2])
+    np.testing.assert_array_equal(final.obs.policy, [[11], [31]])
+    np.testing.assert_array_equal(env.capture_rows[0], [0, 2])
+    np.testing.assert_array_equal(env.capture_episode_steps[0], [3, 3])
+    np.testing.assert_array_equal(env.capture_rewards[0], [11, 31])
+    np.testing.assert_array_equal(env.reset_rows[-1], [0, 1, 2])
+    np.testing.assert_array_equal(state.episode_steps, [0, 0, 0, 1])
+    np.testing.assert_array_equal(state.obs.policy, [[0], [0], [0], [41]])
+    assert env.observation_calls == 2
+    assert not np.shares_memory(final.obs.policy, env.counters)
+    assert not np.shares_memory(final.obs.policy, state.obs.policy)
+    if with_value:
+        np.testing.assert_array_equal(final.obs.value, [[11, -11], [31, -31]])
+        np.testing.assert_array_equal(state.obs.value, [[0, 0], [0, 0], [0, 0], [41, -41]])
+        assert final.obs.value_or_policy is final.obs.value
+        assert not np.shares_memory(final.obs.value, state.obs.value)
+    else:
+        assert final.obs.value is None
+        assert state.obs.value is None
+        assert final.obs.value_or_policy is final.obs.policy
+
+
+@pytest.mark.parametrize("with_value", [False, True])
+def test_retained_final_obs_survives_later_steps_and_init_state(with_value: bool) -> None:
+    env = _FinalCounterEnv(with_value=with_value)
+    env.init_state()
+    env.state.episode_steps[:] = [2, 0, 2, 0]
+    env.counters[:] = [10, 20, 30, 40]
+    retained = _step(env)
+    final = retained.final_obs
+    assert final is not None
+
+    current = _step(env)
+    assert current is not retained
+    assert current.final_obs is None
+    assert retained.final_obs is final
+    assert len(env.capture_rows) == 1
+
+    # A subsequent capture must not overwrite the first transition's snapshot.
+    env.state.episode_steps[1] = 2
+    later = _step(env)
+    assert later.final_obs is not None
+    assert later.final_obs is not final
+    np.testing.assert_array_equal(later.final_obs.env_ids, [1, 3])
+    np.testing.assert_array_equal(later.final_obs.obs.policy, [[23], [43]])
+
+    reinitialized = env.init_state()
+    assert reinitialized.final_obs is None
+    assert retained.final_obs is final
+    np.testing.assert_array_equal(final.env_ids, [0, 2])
+    np.testing.assert_array_equal(final.obs.policy, [[11], [31]])
+    if with_value:
+        np.testing.assert_array_equal(final.obs.value, [[11, -11], [31, -31]])
+    else:
+        assert final.obs.value is None
+
+
+def test_default_final_obs_hook_is_unsupported_on_truncation() -> None:
+    env = _CounterEnv()
+    env.init_state()
+    env.state.episode_steps[2] = 2
+    env.counters[2] = 10.0
+
+    state = _step(env)
+
+    assert state.final_obs is None
+    np.testing.assert_array_equal(state.truncated, [False, False, True, False])
+    np.testing.assert_array_equal(env.reset_rows[-1], [2])
+    np.testing.assert_array_equal(state.obs.policy, [[1], [1], [0], [1]])
+    np.testing.assert_array_equal(state.episode_steps, [1, 1, 0, 1])