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
713from motrix_env_core .direct .env import DirectEnv
814from motrix_env_core .numba .manager .env import ManagerEnv
915from 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
1290def test_np_env_inherits_array_env () -> None :
@@ -21,3 +99,115 @@ def test_manager_env_is_array_env_sibling_of_np_env() -> None:
2199def 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