| from pathlib import Path |
| from types import ModuleType, SimpleNamespace |
| import sys |
|
|
| import numpy as np |
| import torch |
|
|
| from groot.vla.experiment import trex_wandb_video_callback as callback_module |
| from groot.vla.experiment import trex_eval_utils |
| from groot.vla.experiment.trex_track_force_eval_utils import ( |
| TrexTrackForceVisualization, |
| _decode_future_video, |
| overlay_track_motion, |
| ) |
| from groot.vla.experiment.trex_wandb_video_callback import TrexWandbVideoCallback |
|
|
|
|
| def test_track_motion_overlay_outputs_only_three_view_panels(): |
| frames = np.zeros((8, 100, 200, 3), dtype=np.uint8) |
| tracks = np.full((1, 16, 250, 2), 0.5, dtype=np.float32) |
| tracks[0, :, :, 0] += np.linspace(-0.2, 0.2, 16)[:, None] |
| visualization = TrexTrackForceVisualization( |
| video_frames=frames, |
| predicted_track_xy=tracks, |
| target_track_xy=tracks.copy(), |
| target_track_visibility=np.ones((1, 16, 250), dtype=np.float32), |
| anchor_timestamps=np.array([0.0], dtype=np.float64), |
| frame_track_indices=np.array( |
| [2, 4, 6, 8, 10, 12, 14, 15], dtype=np.int64 |
| ), |
| ) |
|
|
| overlaid = overlay_track_motion(visualization, trail_steps=8) |
|
|
| assert overlaid.shape == (8, 50, 300, 3) |
| assert np.count_nonzero(overlaid[:, :, :100]) > 0 |
| assert np.count_nonzero(overlaid[:, :, 100:200]) > 0 |
| assert np.count_nonzero(overlaid[:, :, 200:]) > 0 |
|
|
|
|
| def test_track_motion_overlay_rejects_mismatched_video_indices(): |
| visualization = TrexTrackForceVisualization( |
| video_frames=np.zeros((2, 100, 200, 3), dtype=np.uint8), |
| predicted_track_xy=np.zeros((1, 16, 250, 2), dtype=np.float32), |
| target_track_xy=np.zeros((1, 16, 250, 2), dtype=np.float32), |
| target_track_visibility=np.ones((1, 16, 250), dtype=np.float32), |
| anchor_timestamps=np.array([0.0], dtype=np.float64), |
| frame_track_indices=np.array([1], dtype=np.int64), |
| ) |
|
|
| try: |
| overlay_track_motion(visualization) |
| except ValueError as exc: |
| assert "video and frame-track index counts" in str(exc) |
| else: |
| raise AssertionError("mismatched frame-track indices must fail") |
|
|
|
|
| def test_track_video_starts_with_model_gt_anchor(): |
| conditioning = np.full((250, 2), 0.5, dtype=np.float32) |
| anchored = np.repeat(conditioning[None], 16, axis=0) |
| anchored[:, :, 0] += np.linspace(0.0, 0.1, 16)[:, None] |
| np.testing.assert_allclose(anchored[0], conditioning) |
|
|
| visualization = TrexTrackForceVisualization( |
| video_frames=np.zeros((9, 100, 200, 3), dtype=np.uint8), |
| predicted_track_xy=anchored[None], |
| target_track_xy=anchored[None], |
| target_track_visibility=np.ones((1, 16, 250), dtype=np.float32), |
| anchor_timestamps=np.array([0.0], dtype=np.float64), |
| frame_track_indices=np.array( |
| [-1, 2, 4, 6, 8, 10, 12, 14, 15], dtype=np.int64 |
| ), |
| conditioning_frame=np.zeros((100, 200, 3), dtype=np.uint8), |
| conditioning_track_xy=conditioning[None], |
| conditioning_track_visibility=np.ones( |
| (1, 250), dtype=np.float32 |
| ), |
| ) |
| overlaid = overlay_track_motion(visualization) |
| assert overlaid.shape == (9, 50, 300, 3) |
| assert np.count_nonzero(overlaid[0]) > 0 |
|
|
|
|
| def test_callback_saves_plain_and_track_overlay_videos_separately( |
| tmp_path, monkeypatch |
| ): |
| visualization = TrexTrackForceVisualization( |
| video_frames=np.zeros((8, 100, 200, 3), dtype=np.uint8), |
| predicted_track_xy=np.zeros((1, 16, 250, 2), dtype=np.float32), |
| target_track_xy=np.zeros((1, 16, 250, 2), dtype=np.float32), |
| target_track_visibility=np.ones((1, 16, 250), dtype=np.float32), |
| anchor_timestamps=np.array([0.0], dtype=np.float64), |
| frame_track_indices=np.array( |
| [2, 4, 6, 8, 10, 12, 14, 15], dtype=np.int64 |
| ), |
| ) |
| monkeypatch.setattr( |
| callback_module, |
| "run_trex_track_force_prediction", |
| lambda *args, **kwargs: visualization, |
| ) |
| monkeypatch.setattr( |
| callback_module, |
| "overlay_track_motion", |
| lambda *args, **kwargs: np.full_like( |
| visualization.video_frames, 255 |
| ), |
| ) |
|
|
| written: dict[Path, np.ndarray] = {} |
|
|
| def fake_mimsave(path, frames, **kwargs): |
| written[Path(path)] = np.stack(frames) |
|
|
| monkeypatch.setattr(callback_module.imageio, "mimsave", fake_mimsave) |
|
|
| logged: list[dict] = [] |
| fake_wandb = ModuleType("wandb") |
| fake_wandb.run = object() |
| fake_wandb.define_metric = lambda *args, **kwargs: None |
| fake_wandb.Video = lambda path, **kwargs: {"path": path, **kwargs} |
| fake_wandb.log = logged.append |
| monkeypatch.setitem(sys.modules, "wandb", fake_wandb) |
|
|
| callback = TrexWandbVideoCallback.__new__(TrexWandbVideoCallback) |
| callback.episode_index = 0 |
| callback.num_chunks = 1 |
| callback.every_n_steps = 500 |
| callback.fps = 10 |
| callback.eval_bf16 = False |
| callback.use_dataset_prompt = False |
| callback.prompt = "test task" |
| callback.overlay_tracks = True |
| callback.save_tracks = True |
| callback.track_trail_steps = 8 |
| callback.reconstruction_inference_steps = 1 |
| callback.eval_video_dir = tmp_path / "eval_videos" |
| callback.eval_track_video_dir = tmp_path / "eval_track_videos" |
| callback.eval_track_dir = tmp_path / "eval_tracks" |
| callback._episode = SimpleNamespace( |
| episode_index=0, |
| get_task=lambda row: "dataset task", |
| ) |
| callback._eval_transform = object() |
|
|
| class FakeModel: |
| training = True |
|
|
| def eval(self): |
| self.training = False |
|
|
| def train(self): |
| self.training = True |
|
|
| model = FakeModel() |
| state = SimpleNamespace(is_world_process_zero=True, global_step=500) |
| callback.on_train_begin(None, state, None) |
| callback.on_step_end(None, state, None, model=model) |
|
|
| plain_path = callback.eval_video_dir / "train_step_000500.mp4" |
| track_path = callback.eval_track_video_dir / "train_step_000500.mp4" |
| assert np.count_nonzero(written[plain_path]) == 0 |
| assert np.all(written[track_path] == 255) |
| assert (callback.eval_track_dir / "train_step_000500.npz").is_file() |
| assert TrexWandbVideoCallback.VIDEO_METRIC in logged[0] |
| assert TrexWandbVideoCallback.TRACK_VIDEO_METRIC in logged[0] |
| assert model.training |
|
|
|
|
| def test_episode_frame_reads_are_cached_across_visualization_chunks(monkeypatch): |
| opens = 0 |
|
|
| class FakeCapture: |
| def __init__(self, path): |
| nonlocal opens |
| opens += 1 |
| self.position = 0 |
|
|
| def isOpened(self): |
| return True |
|
|
| def set(self, prop, value): |
| self.position = int(value) |
|
|
| def read(self): |
| value = self.position |
| self.position += 1 |
| return True, np.full((2, 3, 3), value, dtype=np.uint8) |
|
|
| def release(self): |
| pass |
|
|
| monkeypatch.setattr(trex_eval_utils.cv2, "VideoCapture", FakeCapture) |
| episode = trex_eval_utils.TrexEpisode.__new__( |
| trex_eval_utils.TrexEpisode |
| ) |
| episode.length = 10 |
| episode.video_dirs = {"video.head_left": "fake.mp4"} |
| episode._frame_cache = {"video.head_left": {}} |
|
|
| first = episode.get_frames([0, 1, 1, 2], "video.head_left") |
| second = episode.get_frames([1, 2], "video.head_left") |
|
|
| assert opens == 1 |
| assert first[:, 0, 0, 0].tolist() == [0, 1, 1, 2] |
| assert second[:, 0, 0, 0].tolist() == [1, 2] |
|
|
|
|
| def test_track_force_video_decode_reads_tiling_from_policy_config(): |
| seen = {} |
|
|
| class FakeVAE: |
| def decode(self, latents, **kwargs): |
| seen["shape"] = tuple(latents.shape) |
| seen.update(kwargs) |
| return torch.zeros(1, 3, 9, 4, 8) |
|
|
| action_head = SimpleNamespace( |
| config=SimpleNamespace( |
| tiled=False, |
| tile_size_height=20, |
| tile_size_width=21, |
| tile_stride_height=10, |
| tile_stride_width=11, |
| ), |
| vae=FakeVAE(), |
| ) |
| model = SimpleNamespace(action_head=action_head) |
|
|
| frames = _decode_future_video( |
| model, |
| torch.zeros(1, 4, 3, 2, 2), |
| torch.zeros(1, 4, 2, 2, 2), |
| ) |
|
|
| assert frames.shape == (8, 4, 8, 3) |
| assert seen == { |
| "shape": (1, 4, 3, 2, 2), |
| "tiled": False, |
| "tile_size": (20, 21), |
| "tile_stride": (10, 11), |
| } |
|
|
