Skip to content

Commit 18deece

Browse files
committed
fix(fastsac): preserve truncated final observations and stabilize actor
1 parent 14309d3 commit 18deece

24 files changed

Lines changed: 1428 additions & 81 deletions

‎motrix_env_core/src/motrix_env_core/array/env.py‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,18 @@ def value_or_policy(self) -> np.ndarray:
3333
return self.policy if self.value is None else self.value
3434

3535

36+
@dataclass
37+
class FinalObs:
38+
"""Owned pre-reset observations of the current transition's truncated rows.
39+
40+
Rows in ``obs`` align with ``env_ids``. Arrays remain valid across steps;
41+
consumers must treat them as read-only. ``obs.value`` stays optional.
42+
"""
43+
44+
env_ids: np.ndarray
45+
obs: NpObs
46+
47+
3648
@dataclass
3749
class ArrayEnvState:
3850
obs: NpObs
@@ -50,6 +62,8 @@ class ArrayEnvState:
5062
# reflect the latest step: retain values across steps through
5163
# :meth:`process_metrics`, which returns a fresh dict of reduced scalars.
5264
metrics: dict[str, float | np.ndarray] = dataclasses.field(default_factory=dict)
65+
# None when no row truncated or the implementation does not capture finals.
66+
final_obs: FinalObs | None = None
5367

5468
@property
5569
def done(self) -> np.ndarray:
@@ -234,6 +248,15 @@ def reset(self, env_ids: np.ndarray) -> None:
234248
writes have completed.
235249
"""
236250

251+
def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs | None:
252+
"""Capture truncated rows before reset; default implementation is unsupported.
253+
254+
Called after transition and truncation updates, before any reset write.
255+
Implementations must leave normal observations and task state unchanged
256+
and return owned row arrays, without advancing observation history.
257+
"""
258+
return None
259+
237260
def _max_episode_steps(self) -> int | None:
238261
return self._cfg.max_episode_steps
239262

@@ -279,6 +302,8 @@ def physics_step(self):
279302
"""Advance the simulator by the configured substeps."""
280303

281304
def _prev_physics_step(self):
305+
if self._state.final_obs is not None:
306+
self._state = self._state.replace(final_obs=None)
282307
state = self._state
283308
state.reward.fill(0.0)
284309
state.terminated.fill(False)
@@ -302,6 +327,7 @@ def step(self, actions: np.ndarray) -> ArrayEnvState:
302327
303328
Timing: ``apply_action`` -> ``physics_step`` -> ``compute_transition``
304329
(the only full simulator-data refresh) -> episode bookkeeping ->
330+
optional pre-reset ``final_obs`` capture for truncated rows ->
305331
``reset(env_ids)`` for done rows (may refresh selected rows) ->
306332
``compute_observation`` (whole batch, pure construction).
307333
"""
@@ -319,6 +345,10 @@ def step(self, actions: np.ndarray) -> ArrayEnvState:
319345
self._state = self.compute_transition(self._state)
320346
self._state.episode_steps += 1
321347
self._update_truncate()
348+
trunc_ids = np.flatnonzero(self._state.truncated)
349+
if trunc_ids.size:
350+
with self.perf.scope("final_obs"):
351+
self._state = self._state.replace(final_obs=self._capture_final_obs(trunc_ids))
322352
with self.perf.scope("reset"):
323353
self._reset_done_envs()
324354
with self.perf.scope("observation"):
@@ -327,4 +357,4 @@ def step(self, actions: np.ndarray) -> ArrayEnvState:
327357
return self._state
328358

329359

330-
__all__ = ["ArrayEnv", "ArrayEnvState", "EnvCfgType", "NpObs"]
360+
__all__ = ["ArrayEnv", "ArrayEnvState", "EnvCfgType", "FinalObs", "NpObs"]

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

Lines changed: 39 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
import gymnasium as gym
1414
import numpy as np
1515

16-
from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, EnvCfgType, NpObs
16+
from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, EnvCfgType, FinalObs, NpObs
1717
from motrix_env_core.base import EnvCfg, ObsSpace
1818
from motrix_env_core.config import configclass
1919
from motrix_env_core.config.sim_reset import ManagerResetCfg, ResetTermCfg
@@ -285,6 +285,21 @@ class _ObsBufferSlot:
285285
observation: NpObs
286286
kernel_outputs: NumbaKernelOutputs
287287

288+
@classmethod
289+
def bind(cls, observation: NpObs, reward: np.ndarray, terminated: np.ndarray) -> _ObsBufferSlot:
290+
value = observation.value
291+
if value is None:
292+
value = np.empty((observation.policy.shape[0], 0), dtype=observation.policy.dtype)
293+
return cls(
294+
observation=observation,
295+
kernel_outputs=NumbaKernelOutputs(
296+
policy_obs=observation.policy,
297+
value_obs=value,
298+
reward=reward,
299+
terminated=terminated,
300+
),
301+
)
302+
288303

289304
@dataclass
290305
class SwapObsBuffer:
@@ -311,21 +326,7 @@ def create(
311326
reward: np.ndarray,
312327
terminated: np.ndarray,
313328
) -> SwapObsBuffer:
314-
def make_slot(observation: NpObs) -> _ObsBufferSlot:
315-
value = observation.value
316-
if value is None:
317-
value = np.empty((observation.policy.shape[0], 0), dtype=observation.policy.dtype)
318-
return _ObsBufferSlot(
319-
observation=observation,
320-
kernel_outputs=NumbaKernelOutputs(
321-
policy_obs=observation.policy,
322-
value_obs=value,
323-
reward=reward,
324-
terminated=terminated,
325-
),
326-
)
327-
328-
return cls((make_slot(first), make_slot(second)))
329+
return cls((_ObsBufferSlot.bind(first, reward, terminated), _ObsBufferSlot.bind(second, reward, terminated)))
329330

330331
@property
331332
def current(self) -> _ObsBufferSlot:
@@ -370,6 +371,7 @@ def __init__(self, cfg: EnvCfgType, num_envs: int = 1, backend: str | None = Non
370371
self.metrics: dict[str, np.ndarray] = {}
371372
self._kernel_outputs: NumbaKernelOutputs | None = None
372373
self._observation_buffer: SwapObsBuffer | None = None
374+
self._final_obs_slot: _ObsBufferSlot | None = None
373375
# Persistent kernel-bound sim-reset-request buffer
374376
# (``ctx.sim_reset_requested``), cleared before each physics step and
375377
# consumed by _reset_done_envs.
@@ -509,6 +511,8 @@ def init_state(self) -> ArrayEnvState:
509511
state.terminated,
510512
)
511513
self._kernel_outputs = self._observation_buffer.current.kernel_outputs
514+
final_obs = self._zeros_obs(self.observation_space)
515+
self._final_obs_slot = _ObsBufferSlot.bind(final_obs, state.reward, state.terminated)
512516
self._refresh_sim_reads()
513517
inputs = self._kernel_inputs
514518
self._validate_kernel_context(inputs)
@@ -548,6 +552,25 @@ def _execute_observe_kernel(self, inputs: tuple[Any, ...]) -> None:
548552
assert self._kernel_outputs is not None
549553
self._task_program.observe_kernel(inputs, self._kernel_outputs)
550554

555+
def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs:
556+
"""Observe pre-reset inputs after command on_transition.
557+
558+
The compiled observe pass is batch-wide. Dedicated outputs preserve the
559+
normal swap buffer; advanced indexing creates owned compact snapshots.
560+
"""
561+
assert self._task_program is not None
562+
assert self._final_obs_slot is not None
563+
slot = self._final_obs_slot
564+
self._task_program.observe_kernel(self._kernel_inputs, slot.kernel_outputs)
565+
obs = slot.observation
566+
return FinalObs(
567+
env_ids=env_ids,
568+
obs=NpObs(
569+
policy=obs.policy[env_ids],
570+
value=None if obs.value is None else obs.value[env_ids],
571+
),
572+
)
573+
551574
def _execute_task_kernel(self, inputs: tuple[Any, ...]) -> None:
552575
"""Evaluate and observe without advancing host-side commands."""
553576
self._execute_evaluate_kernel(inputs)

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,12 @@ def __call__(self, ctx: BuildContext) -> ObsTerm:
3030

3131
@dataclass(frozen=True, slots=True, init=False)
3232
class ObsTerm(BaseTerm):
33-
"""Observation dispatch term with a fixed output width."""
33+
"""Pure observation dispatch with a fixed output width.
34+
35+
Read context and write only the supplied output, except observation noise
36+
may consume ctx.rand. Do not advance history or commands: truncated steps
37+
may observe again before reset, with the RNG restored after that snapshot.
38+
"""
3439

3540
size: int
3641

‎motrix_env_core/tests/test_array_env.py‎

Lines changed: 192 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,90 @@
11
# Copyright Motphys Technology Co., Ltd. 2025, 2026
22
# SPDX-License-Identifier: Apache-2.0
33

4-
"""Regression tests for the ArrayEnv frontend hierarchy (issue #222 MR1)."""
4+
"""Regression tests for the ArrayEnv hierarchy and final-observation lifecycle."""
55

6-
from motrix_env_core.array.env import ArrayEnv
6+
import gymnasium as gym
7+
import numpy as np
8+
import pytest
9+
10+
from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, FinalObs, NpObs
11+
from motrix_env_core.base import EnvCfg, ObsSpace
12+
from motrix_env_core.config.scene import SceneCfg
713
from motrix_env_core.direct.env import DirectEnv
814
from motrix_env_core.numba.manager.env import ManagerEnv
915
from motrix_env_core.registry import _infer_sim_backend
16+
from motrix_env_core.sim.backend import RenderConfig, SimRenderer
17+
18+
19+
class _CounterEnv(ArrayEnv):
20+
"""Exercise the frontend lifecycle without constructing a simulator backend."""
21+
22+
def __init__(self, *, with_value: bool = False):
23+
super().__init__(EnvCfg(scene=SceneCfg(), ctrl_dt=1.0, max_episode_seconds=3.0), num_envs=4)
24+
policy_space = gym.spaces.Box(-np.inf, np.inf, shape=(1,), dtype=np.float32)
25+
value_space = gym.spaces.Box(-np.inf, np.inf, shape=(2,), dtype=np.float32) if with_value else None
26+
self._observation_space = ObsSpace(policy=policy_space, value=value_space)
27+
self._action_space = gym.spaces.Box(-1.0, 1.0, shape=(1,), dtype=np.float32)
28+
self.counters = np.zeros(self.num_envs, dtype=np.float32)
29+
self.next_terminated = np.zeros(self.num_envs, dtype=bool)
30+
self.reset_rows: list[np.ndarray] = []
31+
self.physics_steps = 0
32+
self.observation_calls = 0
33+
34+
@property
35+
def observation_space(self) -> ObsSpace:
36+
return self._observation_space
37+
38+
@property
39+
def action_space(self) -> gym.spaces.Box:
40+
return self._action_space
41+
42+
def reset(self, env_ids: np.ndarray) -> None:
43+
self.reset_rows.append(env_ids.copy())
44+
self.counters[env_ids] = 0.0
45+
46+
def apply_action(self, actions: np.ndarray, state: ArrayEnvState) -> ArrayEnvState:
47+
return state
48+
49+
def physics_step(self) -> None:
50+
self.physics_steps += 1
51+
self.counters += 1.0
52+
53+
def compute_transition(self, state: ArrayEnvState) -> ArrayEnvState:
54+
state.reward[:] = self.counters
55+
state.terminated[:] = self.next_terminated
56+
return state
57+
58+
def compute_observation(self, state: ArrayEnvState) -> ArrayEnvState:
59+
self.observation_calls += 1
60+
state.obs.policy[:, 0] = self.counters
61+
if state.obs.value is not None:
62+
state.obs.value[:, 0] = self.counters
63+
state.obs.value[:, 1] = -self.counters
64+
return state
65+
66+
def create_renderer(self, config: RenderConfig) -> SimRenderer:
67+
raise NotImplementedError("The counter environment has no simulator or renderer")
68+
69+
70+
class _FinalCounterEnv(_CounterEnv):
71+
def __init__(self, *, with_value: bool = False):
72+
super().__init__(with_value=with_value)
73+
self.capture_rows: list[np.ndarray] = []
74+
self.capture_episode_steps: list[np.ndarray] = []
75+
self.capture_rewards: list[np.ndarray] = []
76+
77+
def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs:
78+
self.capture_rows.append(env_ids.copy())
79+
self.capture_episode_steps.append(self.state.episode_steps[env_ids].copy())
80+
self.capture_rewards.append(self.state.reward[env_ids].copy())
81+
policy = self.counters[env_ids, None].copy()
82+
value = np.concatenate((policy, -policy), axis=1) if self.has_value_observation else None
83+
return FinalObs(env_ids=env_ids.copy(), obs=NpObs(policy=policy, value=value))
84+
85+
86+
def _step(env: _CounterEnv) -> ArrayEnvState:
87+
return env.step(np.zeros((env.num_envs, 1), dtype=np.float32))
1088

1189

1290
def test_np_env_inherits_array_env() -> None:
@@ -21,3 +99,115 @@ def test_manager_env_is_array_env_sibling_of_np_env() -> None:
2199
def test_registry_infers_np_backend_for_manager_env() -> None:
22100
assert _infer_sim_backend(ManagerEnv) == "np"
23101
assert _infer_sim_backend(DirectEnv) == "np"
102+
103+
104+
@pytest.mark.parametrize("with_value", [False, True])
105+
def test_final_obs_absent_on_init_and_nontruncated_step(with_value: bool) -> None:
106+
env = _FinalCounterEnv(with_value=with_value)
107+
initial = env.init_state()
108+
109+
assert initial.final_obs is None
110+
np.testing.assert_array_equal(initial.obs.policy, np.zeros((4, 1)))
111+
np.testing.assert_array_equal(env.reset_rows[0], np.arange(4))
112+
assert env.physics_steps == 0
113+
assert env.observation_calls == 1
114+
assert env.capture_rows == []
115+
116+
state = _step(env)
117+
118+
assert state.final_obs is None
119+
assert not state.truncated.any()
120+
assert not state.terminated.any()
121+
np.testing.assert_array_equal(state.episode_steps, np.ones(4))
122+
np.testing.assert_array_equal(state.obs.policy, np.ones((4, 1)))
123+
assert env.physics_steps == 1
124+
assert env.observation_calls == 2
125+
assert len(env.reset_rows) == 1
126+
assert env.capture_rows == []
127+
128+
129+
@pytest.mark.parametrize("with_value", [False, True])
130+
def test_final_obs_contains_only_truncated_rows_before_auto_reset(with_value: bool) -> None:
131+
env = _FinalCounterEnv(with_value=with_value)
132+
env.init_state()
133+
env.state.episode_steps[:] = [2, 0, 2, 0]
134+
env.counters[:] = [10, 20, 30, 40]
135+
env.next_terminated[1] = True
136+
137+
state = _step(env)
138+
139+
final = state.final_obs
140+
assert isinstance(final, FinalObs)
141+
assert isinstance(final.obs, NpObs)
142+
np.testing.assert_array_equal(state.truncated, [True, False, True, False])
143+
np.testing.assert_array_equal(state.terminated, [False, True, False, False])
144+
np.testing.assert_array_equal(final.env_ids, [0, 2])
145+
np.testing.assert_array_equal(final.obs.policy, [[11], [31]])
146+
np.testing.assert_array_equal(env.capture_rows[0], [0, 2])
147+
np.testing.assert_array_equal(env.capture_episode_steps[0], [3, 3])
148+
np.testing.assert_array_equal(env.capture_rewards[0], [11, 31])
149+
np.testing.assert_array_equal(env.reset_rows[-1], [0, 1, 2])
150+
np.testing.assert_array_equal(state.episode_steps, [0, 0, 0, 1])
151+
np.testing.assert_array_equal(state.obs.policy, [[0], [0], [0], [41]])
152+
assert env.observation_calls == 2
153+
assert not np.shares_memory(final.obs.policy, env.counters)
154+
assert not np.shares_memory(final.obs.policy, state.obs.policy)
155+
if with_value:
156+
np.testing.assert_array_equal(final.obs.value, [[11, -11], [31, -31]])
157+
np.testing.assert_array_equal(state.obs.value, [[0, 0], [0, 0], [0, 0], [41, -41]])
158+
assert final.obs.value_or_policy is final.obs.value
159+
assert not np.shares_memory(final.obs.value, state.obs.value)
160+
else:
161+
assert final.obs.value is None
162+
assert state.obs.value is None
163+
assert final.obs.value_or_policy is final.obs.policy
164+
165+
166+
@pytest.mark.parametrize("with_value", [False, True])
167+
def test_retained_final_obs_survives_later_steps_and_init_state(with_value: bool) -> None:
168+
env = _FinalCounterEnv(with_value=with_value)
169+
env.init_state()
170+
env.state.episode_steps[:] = [2, 0, 2, 0]
171+
env.counters[:] = [10, 20, 30, 40]
172+
retained = _step(env)
173+
final = retained.final_obs
174+
assert final is not None
175+
176+
current = _step(env)
177+
assert current is not retained
178+
assert current.final_obs is None
179+
assert retained.final_obs is final
180+
assert len(env.capture_rows) == 1
181+
182+
# A subsequent capture must not overwrite the first transition's snapshot.
183+
env.state.episode_steps[1] = 2
184+
later = _step(env)
185+
assert later.final_obs is not None
186+
assert later.final_obs is not final
187+
np.testing.assert_array_equal(later.final_obs.env_ids, [1, 3])
188+
np.testing.assert_array_equal(later.final_obs.obs.policy, [[23], [43]])
189+
190+
reinitialized = env.init_state()
191+
assert reinitialized.final_obs is None
192+
assert retained.final_obs is final
193+
np.testing.assert_array_equal(final.env_ids, [0, 2])
194+
np.testing.assert_array_equal(final.obs.policy, [[11], [31]])
195+
if with_value:
196+
np.testing.assert_array_equal(final.obs.value, [[11, -11], [31, -31]])
197+
else:
198+
assert final.obs.value is None
199+
200+
201+
def test_default_final_obs_hook_is_unsupported_on_truncation() -> None:
202+
env = _CounterEnv()
203+
env.init_state()
204+
env.state.episode_steps[2] = 2
205+
env.counters[2] = 10.0
206+
207+
state = _step(env)
208+
209+
assert state.final_obs is None
210+
np.testing.assert_array_equal(state.truncated, [False, False, True, False])
211+
np.testing.assert_array_equal(env.reset_rows[-1], [2])
212+
np.testing.assert_array_equal(state.obs.policy, [[1], [1], [0], [1]])
213+
np.testing.assert_array_equal(state.episode_steps, [1, 1, 0, 1])

0 commit comments

Comments
 (0)