Let task owners inject the ManagerBasedRlEnv class and keep pre-advance clip ownership in the tensor sampler

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

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,

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

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()