From e175a6be5f796be969b992e994f70b2f9ec02867 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Fri, 9 Oct 2026 19:31:53 +0800 Subject: [PATCH] fix(motion): gather reference frames from device carrier in full refresh TensorMotionCommand._refresh_motion_torch read the host NumPy MotionSampler.current_frames on the full (rows=None) refresh path, but c9a8169d removed the host mirror sync when stepping moved to TensorMotionSampler.step_full, leaving the host frames frozen at zero. Every per-step refresh re-ingested frame 0 of the reference clip into the command, motion-anchor observations, and tracking rewards for all TensorMotionCommand owners (flashsac all sims, sac/mjwarp), in both training and eval. Gather from the device time_steps carrier instead, matching the selected-rows branch, and add a regression test that keeps the host sampler frozen at frame 0 and asserts the full refresh gathers the advanced device frames. Fixes #2104 --- .../motion_tracking/common/manager_terms.py | 6 ++- tests/tasks/test_motion_term_parity.py | 53 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/src/unilab/tasks/motion_tracking/common/manager_terms.py b/src/unilab/tasks/motion_tracking/common/manager_terms.py index 519498bb7..11a3168b7 100644 --- a/src/unilab/tasks/motion_tracking/common/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/common/manager_terms.py @@ -1483,7 +1483,11 @@ def _refresh_motion_torch(self, rows: torch.Tensor | None = None) -> None: """Gather motion rows and publish the device-resident command carrier.""" frames: np.ndarray | torch.Tensor if rows is None: - frames = self.sampler.current_frames + frames = ( + cast(torch.Tensor, self.time_steps) + if isinstance(self.time_steps, torch.Tensor) + else self.sampler.current_frames + ) elif isinstance(self.time_steps, torch.Tensor): frames = cast(torch.Tensor, self.time_steps)[rows] else: diff --git a/tests/tasks/test_motion_term_parity.py b/tests/tasks/test_motion_term_parity.py index ddbea7eff..f7fde9519 100644 --- a/tests/tasks/test_motion_term_parity.py +++ b/tests/tasks/test_motion_term_parity.py @@ -969,6 +969,59 @@ def test_tensor_command_publishes_sampler_advance_exactly_once() -> None: torch.testing.assert_close(command.time_steps, torch.tensor([5, 1, 10], dtype=torch.int32)) +def test_tensor_command_full_refresh_gathers_advanced_frames() -> None: + command = mt.TensorMotionCommand.__new__(mt.TensorMotionCommand) + command._device = torch.device("cpu") + command.time_steps = torch.tensor([4, 0, 9], dtype=torch.int32) + command.tensor_sampler = mt.TensorMotionSampler( + mode="adaptive", + num_envs=3, + num_frames=11, + clip_offsets=np.asarray([0], dtype=np.int64), + clip_end_frames=np.asarray([100], dtype=np.int32), + bin_count=1, + adaptive_lambda=0.8, + adaptive_kernel_size=1, + adaptive_uniform_ratio=0.1, + adaptive_alpha=0.001, + start_ratio=0.0, + initial_frames=np.asarray([4, 0, 9], dtype=np.int32), + initial_clip_end_frames=np.asarray([100, 100, 100], dtype=np.int32), + device=torch.device("cpu"), + ) + command.tensor_sampler.current_frames.copy_(command.time_steps) + command.current_clip_end_frames = torch.tensor([90, 90, 90], dtype=torch.int32) + command._clip_offsets_torch = torch.tensor([0], dtype=torch.int64) + command._clip_end_frames_torch = torch.tensor([100], dtype=torch.int64) + command._tensor_all_rows = torch.arange(3, dtype=torch.int64) + command.cfg = SimpleNamespace(params=SimpleNamespace(truncate_on_clip_end=True)) + command._tensor_post_compute_env_ids = None + command._resample_ingested_ids = None + command._tensor_resample_ingested = None + command._env = SimpleNamespace( + termination_manager=SimpleNamespace(terminated=torch.tensor([False, False, False])), + reset_buf=torch.tensor([False, False, False]), + ) + # The host NumPy sampler is never stepped on the tensor path; keep it frozen + # at frame 0 so the refresh must gather from the device carrier instead. + command.sampler = SimpleNamespace(current_frames=np.zeros(3, dtype=np.int64)) + + captured: dict[str, Any] = {} + command._motion_packet = lambda frames: frames + + def _capture_ingest(rows: torch.Tensor, packet: torch.Tensor) -> None: + captured["rows"] = rows + captured["frames"] = packet + + command._ingest_motion_packet = _capture_ingest + command._update_command(None) + + torch.testing.assert_close(captured["rows"], torch.arange(3, dtype=torch.int64)) + torch.testing.assert_close( + captured["frames"].to(torch.int64), torch.tensor([5, 1, 10], dtype=torch.int64) + ) + + def test_tensor_command_syncs_sampler_mirrors_on_selected_rows_only() -> None: command = mt.TensorMotionCommand.__new__(mt.TensorMotionCommand) command._device = torch.device("cpu")