From d8afa17d54f9b07c714ea15b312c1f426a181698 Mon Sep 17 00:00:00 2001 From: caozx1110 Date: Sat, 24 Oct 2026 20:20:13 +0800 Subject: [PATCH] fix: preserve SONIC factory and tensor clip boundary contracts Port the task-owned ManagerBasedRlEnv factory extension and pre-advance clip ownership semantics from the SONIC fork cded6de532ed854d47bab0f65cb012d4730d44cd onto UniLab v1.5.1 (13ba6bb61e14). Preserve the full-width bool return ABI and stable sampler scratch buffers. Source: sonic_unilab/docs/patches/unilab_151_sonic_tensor_contract.patch; upstream code remains under its existing license. Add CPU boundary/lifecycle regressions. Validation: 158 passed, 6 skipped, 3 Go2-asset-dependent cases deselected; Ruff, mypy and pyright passed. No GPU validation. Refs unilabsim/sonic_unilab#30 --- src/unilab/envs/manager_based_rl_env.py | 14 +++++- .../motion_tracking/common/tensor_sampler.py | 15 +++--- tests/envs/test_manager_based_rl_env.py | 43 +++++++++++++++++ tests/tasks/test_tensor_motion_sampler.py | 47 ++++++++++++++++++- 4 files changed, 109 insertions(+), 10 deletions(-) diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index a2ecb8036..9cbea455e 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -1928,12 +1928,22 @@ def make_manager_based_rl_env( cfg: ManagerBasedRlEnvCfg, num_envs: int = 1, backend_type: str = "mujoco", + *, + env_cls: type[ManagerBasedRlEnv] | None = None, ) -> ManagerBasedRlEnv: - """Construct the generic Registry-owned Manager-Based production runtime.""" + """Construct the shared runtime, optionally using a task-owned subclass. + + The subclass shares the backend allocation and cleanup lifecycle and must + accept the same (cfg, backend, num_envs) constructor contract. + """ from unilab.base.registry import ( _SUPPORTED_SIM_BACKENDS, ) + if env_cls is not None and ( + not isinstance(env_cls, type) or not issubclass(env_cls, ManagerBasedRlEnv) + ): + raise TypeError("env_cls must be a ManagerBasedRlEnv subclass") if not isinstance(cfg, ManagerBasedRlEnvCfg): raise TypeError( "make_manager_based_rl_env expected ManagerBasedRlEnvCfg, " @@ -1982,7 +1992,7 @@ def make_manager_based_rl_env( require_fixed_variant_support( backend.get_dr_capabilities(), cfg.scene.fixed_variant_plan ) - return ManagerBasedRlEnv(cfg, backend, num_envs) + return (ManagerBasedRlEnv if env_cls is None else env_cls)(cfg, backend, num_envs) except Exception: backend.cleanup_scene_assets() raise diff --git a/src/unilab/tasks/motion_tracking/common/tensor_sampler.py b/src/unilab/tasks/motion_tracking/common/tensor_sampler.py index 1a2adfa83..865386bc3 100644 --- a/src/unilab/tasks/motion_tracking/common/tensor_sampler.py +++ b/src/unilab/tasks/motion_tracking/common/tensor_sampler.py @@ -342,17 +342,18 @@ def step_full(self, active: torch.Tensor, time_steps: torch.Tensor) -> torch.Ten or time_steps.device != self.device ): raise ValueError("MotionSampler time_steps must be int32 device frames") + # Resolve ownership before advancing so a terminal frame cannot enter + # the next clip. Keep a valid reference until the owner handles done. + clip_ends = self._step_clip_ends + torch.index_select(self._frame_clip_ends, 0, self.current_frames, out=clip_ends) next_frames = self._step_next_frames torch.add(self.current_frames, active, out=next_frames) - next_frames.clamp_(max=self.num_frames - 1) - frames = self._step_clamped_rows - frames.copy_(next_frames) - clip_ends = self._step_clip_ends - torch.index_select(self._frame_clip_ends, 0, frames, out=clip_ends) done = self._step_done - torch.gt(frames, clip_ends, out=done) + torch.gt(next_frames, clip_ends, out=done) done.logical_and_(active) - self.current_frames.copy_(next_frames) + frames = self._step_clamped_rows + torch.minimum(next_frames, clip_ends, out=frames) + self.current_frames.copy_(frames) torch.where( active, clip_ends, diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 816790964..9c62f6539 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -4846,3 +4846,46 @@ def test_env_reset_backfills_history_terms_without_touching_other_rows() -> None torch.testing.assert_close(obs["obs"][1].reshape(3, 2), reset_slots) torch.testing.assert_close(reset_slots[-1], torch.zeros(2)) env.close() + + +def test_factory_accepts_explicit_task_subclass(monkeypatch: pytest.MonkeyPatch) -> None: + cfg = _make_cfg(include_optional_managers=False) + assert cfg.scene is not None + cfg.scene.entities["robot"] = EntityCfg(root_body_name="base", actuator_names=("motor",)) + backend = _FakeBackend(2) + monkeypatch.setattr(manager_env_module, "create_backend", lambda *a, **k: backend) + + class TaskEnv(ManagerBasedRlEnv): + def __init__(self, received_cfg, received_backend, received_num_envs): + self.received = (received_cfg, received_backend, received_num_envs) + + result = make_manager_based_rl_env(cfg, num_envs=2, env_cls=TaskEnv) + assert type(result) is TaskEnv + assert result.received == (cfg, backend, 2) + assert backend.cleanup_calls == 0 + + +@pytest.mark.parametrize("env_cls", [object, lambda *a: None, 42]) +def test_factory_rejects_invalid_task_class_before_allocation(monkeypatch, env_cls) -> None: + def reject_backend(*args, **kwargs): + raise AssertionError("invalid owner must not allocate a backend") + + monkeypatch.setattr(manager_env_module, "create_backend", reject_backend) + with pytest.raises(TypeError, match="env_cls"): + make_manager_based_rl_env(_make_cfg(include_optional_managers=False), env_cls=env_cls) + + +def test_factory_cleans_backend_when_task_subclass_fails(monkeypatch) -> None: + cfg = _make_cfg(include_optional_managers=False) + assert cfg.scene is not None + cfg.scene.entities["robot"] = EntityCfg(root_body_name="base", actuator_names=("motor",)) + backend = _FakeBackend(2) + monkeypatch.setattr(manager_env_module, "create_backend", lambda *a, **k: backend) + + class FailingTaskEnv(ManagerBasedRlEnv): + def __init__(self, cfg, backend, num_envs): + raise RuntimeError("task construction failed") + + with pytest.raises(RuntimeError, match="task construction failed"): + make_manager_based_rl_env(cfg, num_envs=2, env_cls=FailingTaskEnv) + assert backend.cleanup_calls == 1 diff --git a/tests/tasks/test_tensor_motion_sampler.py b/tests/tasks/test_tensor_motion_sampler.py index 2dad82431..e82b996db 100644 --- a/tests/tasks/test_tensor_motion_sampler.py +++ b/tests/tasks/test_tensor_motion_sampler.py @@ -98,7 +98,7 @@ def test_step_full_advances_active_rows_and_reports_done() -> None: # Repeated applies keep advancing active rows exactly once per call. second = sampler.step_full(active, time_steps) - assert second.tolist() == [False, False, False, False] + assert second.tolist() == [True, False, False, False] torch.testing.assert_close( sampler.current_frames, torch.tensor([31, 10, 13, 3], dtype=torch.int32) ) @@ -305,3 +305,48 @@ def test_frame_indexed_clip_ends_match_searchsorted_ownership() -> None: expected = sampler._clip_end_frames.index_select(0, expected_indices) torch.testing.assert_close(sampler._frame_clip_ends, expected) + + +@pytest.mark.parametrize( + ("frame", "active", "expected", "done"), + [ + (1, True, 2, False), + (2, True, 2, True), + (4, True, 5, False), + (5, True, 5, True), + (2, False, 2, False), + (5, False, 5, False), + ], +) +def test_step_full_preserves_original_clip_ownership(frame, active, expected, done): + sampler = MotionSampler( + mode="mixed", + num_envs=1, + num_frames=6, + clip_offsets=torch.tensor([0, 3], dtype=torch.int64), + clip_end_frames=torch.tensor([2, 5], dtype=torch.int32), + bin_count=2, + adaptive_lambda=0.8, + adaptive_kernel_size=1, + adaptive_uniform_ratio=0.1, + adaptive_alpha=0.5, + start_ratio=0.0, + initial_frames=torch.tensor([frame], dtype=torch.int32), + initial_clip_end_frames=torch.tensor([2 if frame < 3 else 5], dtype=torch.int32), + device=torch.device("cpu"), + ) + end_before = sampler.current_clip_end_frames.clone() + scratch = ( + sampler._step_next_frames, + sampler._step_clamped_rows, + sampler._step_clip_ends, + sampler._step_done, + ) + pointers = [value.data_ptr() for value in scratch] + times = torch.full((1,), -1, dtype=torch.int32) + result = sampler.step_full(torch.tensor([active]), times) + assert result.dtype == torch.bool and result.tolist() == [done] + assert sampler.current_frames.tolist() == times.tolist() == [expected] + torch.testing.assert_close(sampler.current_clip_end_frames, end_before, rtol=0, atol=0) + assert [value.data_ptr() for value in scratch] == pointers + assert result.data_ptr() == sampler._step_done.data_ptr()