#!/usr/bin/env python3 """Standalone T-Rex inference (no websocket server). Loads your LoRA checkpoint, runs causal video+action diffusion on dataset frames, and saves the *predicted* future video as MP4. Usage: CUDA_VISIBLE_DEVICES=0 python scripts/inference_trex.py \ --model_path /scratch1/home/zhicao/dreamzero/checkpoints/dreamzero_trex_wan22_lora/checkpoint-5000 \ --dataset_path /scratch1/home/zhicao/dreamzero/data/trex_small \ --episode 0 \ --num_chunks 8 \ --output_dir results_trex_infer Why not server/client? The official server (socket_test_optimized_AR.py) + test_client_AR.py split model hosting from the robot/eval loop for multi-GPU distributed inference. For offline debugging you only need GrootSimPolicy + lazy_joint_forward_causal. """ from __future__ import annotations import argparse import glob import os import time import cv2 import imageio import numpy as np import pyarrow.parquet as pq import torch import torch._dynamo import torch.distributed as dist from einops import rearrange from tianshou.data import Batch torch._dynamo.config.disable = True from groot.vla.data.schema import EmbodimentTag from groot.vla.data.transform import ComposedModalityTransform from groot.vla.model.n1_5.sim_policy import GrootSimPolicy # Modality keys (must match transform_trex / modality_config_trex) VIDEO_KEYS = [ "video.head_left", "video.left_wrist", "video.right_wrist", ] STATE_KEYS = { "state.left_arm": (0, 7), "state.left_hand": (7, 29), "state.right_arm": (29, 36), "state.right_hand": (36, 58), } VIDEO_FOLDERS = { "video.head_left": "observation.images.head_left", "video.left_wrist": "observation.images.left_wrist", "video.right_wrist": "observation.images.right_wrist", } # Causal chunk schedule (same idea as test_client_AR.py / DROID server) RELATIVE_OFFSETS = [-23, -16, -8, 0] ACTION_HORIZON = 24 def get_expected_video_resolution(policy: GrootSimPolicy) -> tuple[int, int]: """Return (height, width) that eval_transform VideoToTensor expects.""" cfg = policy.trained_model.action_head.config target_h = getattr(cfg, "target_video_height", None) target_w = getattr(cfg, "target_video_width", None) if target_h is not None and target_w is not None: return int(target_h), int(target_w) eval_transform = getattr(policy, "eval_transform", None) if isinstance(eval_transform, ComposedModalityTransform): for t in eval_transform.transforms: res = getattr(t, "original_resolutions", None) if res: w, h = next(iter(res.values())) return int(h), int(w) return 160, 320 def resize_frames(frames: np.ndarray, target_h: int, target_w: int) -> np.ndarray: """Resize (H,W,C) or (T,H,W,C) uint8 frames to (target_h, target_w).""" if frames.ndim == 3: if (frames.shape[0], frames.shape[1]) == (target_h, target_w): return frames return cv2.resize(frames, (target_w, target_h), interpolation=cv2.INTER_LINEAR) return np.stack( [cv2.resize(f, (target_w, target_h), interpolation=cv2.INTER_LINEAR) for f in frames], axis=0, ) class TrexEpisode: """One T-Rex episode from LeRobot v2 layout.""" def __init__(self, dataset_root: str, episode_index: int): pq_path = os.path.join( dataset_root, "data", f"chunk-{episode_index // 1000:03d}", f"episode_{episode_index:06d}.parquet", ) if not os.path.isfile(pq_path): raise FileNotFoundError(pq_path) self.table = pq.read_table(pq_path) self.length = self.table.num_rows self.episode_index = episode_index self.root = dataset_root self.video_dirs = {} for key, folder in VIDEO_FOLDERS.items(): pattern = os.path.join( dataset_root, "videos", "**", folder, f"episode_{episode_index:06d}.mp4", ) hits = sorted(glob.glob(pattern, recursive=True)) if not hits: raise FileNotFoundError(f"No video for {key}: {pattern}") self.video_dirs[key] = hits[0] print( f"TrexEpisode {episode_index}: {self.length} steps, " f"{len(self.video_dirs)} cameras" ) def get_task(self, row: int) -> str: try: return str(self.table.column("annotation.task")[row].as_py()) except Exception: return "" def get_state(self, row: int) -> np.ndarray: return np.array(self.table.column("observation.state")[row].as_py(), dtype=np.float64) def get_frame(self, row: int, video_key: str) -> np.ndarray: cap = cv2.VideoCapture(self.video_dirs[video_key]) cap.set(cv2.CAP_PROP_POS_FRAMES, row) ok, frame = cap.read() cap.release() if not ok: raise RuntimeError(f"Failed frame {row} from {self.video_dirs[video_key]}") return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) def build_chunk_schedule(total_frames: int, num_chunks: int) -> list[list[int]]: """First chunk: 1 frame; later chunks: 4 frames ending at anchor.""" chunks: list[list[int]] = [] anchor = 0 for i in range(num_chunks): if i == 0: indices = [0] else: indices = [max(anchor + off, 0) for off in RELATIVE_OFFSETS] if indices[-1] >= total_frames: break chunks.append(indices) anchor += ACTION_HORIZON return chunks def build_obs( episode: TrexEpisode, frame_indices: list[int], prompt: str, video_height: int, video_width: int, ) -> dict: obs: dict = {} anchor = frame_indices[-1] state = episode.get_state(anchor) for key in VIDEO_KEYS: frames = np.stack([episode.get_frame(i, key) for i in frame_indices], axis=0) frames = resize_frames(frames.astype(np.uint8), video_height, video_width) if len(frame_indices) == 1: obs[key] = frames[0] # (H, W, 3) else: obs[key] = frames # (T, H, W, 3) for key, (s, e) in STATE_KEYS.items(): obs[key] = state[s:e].reshape(1, -1).astype(np.float64) obs["annotation.task"] = prompt return obs def decode_video_latents(policy: GrootSimPolicy, video_chunks: list[torch.Tensor]) -> np.ndarray: """Concat latent chunks along time, VAE decode -> (T, H, W, 3) uint8.""" if not video_chunks: raise ValueError("No video chunks to decode") cat = torch.cat(video_chunks, dim=2) ah = policy.trained_model.action_head frames = ah.vae.decode( cat, tiled=ah.tiled, tile_size=(ah.tile_size_height, ah.tile_size_width), tile_stride=(ah.tile_stride_height, ah.tile_stride_width), ) frames = rearrange(frames, "B C T H W -> B T H W C")[0] frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) return frames def save_mp4(path: str, frames: np.ndarray, fps: int = 5) -> None: os.makedirs(os.path.dirname(path) or ".", exist_ok=True) imageio.mimsave(path, list(frames), fps=fps, codec="libx264") print(f"Saved {len(frames)} frames -> {path}") def run(args: argparse.Namespace) -> None: if not dist.is_initialized(): os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29500") dist.init_process_group(backend="gloo", world_size=1, rank=0) print(f"Loading checkpoint: {args.model_path}") policy = GrootSimPolicy( embodiment_tag=EmbodimentTag.TREX, model_path=args.model_path, device=args.device, ) print("Model loaded.") video_height, video_width = get_expected_video_resolution(policy) print(f"Resizing input video to {video_height}x{video_width} (HxW) for eval_transform") episode = TrexEpisode(args.dataset_path, args.episode) prompt = episode.get_task(0) if args.use_dataset_prompt else args.prompt schedule = build_chunk_schedule(episode.length, args.num_chunks) print(f"Prompt: {prompt!r}") print(f"Running {len(schedule)} causal chunks: {schedule}") os.makedirs(args.output_dir, exist_ok=True) video_chunks: list[torch.Tensor] = [] times = [] for ci, frame_indices in enumerate(schedule): obs = build_obs(episode, frame_indices, prompt, video_height, video_width) t0 = time.perf_counter() with torch.inference_mode(): result, video_pred = policy.lazy_joint_forward_causal(Batch(obs=obs)) elapsed = time.perf_counter() - t0 times.append(elapsed) video_chunks.append(video_pred) act = result.act print( f" chunk {ci:02d} frames={frame_indices} " f"infer={elapsed:.2f}s start_frame={policy.trained_model.action_head.current_start_frame}" ) if ci == 0: print(f" action keys: {[k for k in dir(act) if k.startswith('action.')]}") pred_frames = decode_video_latents(policy, video_chunks) out_pred = os.path.join( args.output_dir, f"ep{args.episode:06d}_pred.mp4", ) save_mp4(out_pred, pred_frames, fps=args.fps) if args.save_input_clip: # Save the conditioning frames (head camera) for reference input_frames = [] max_idx = min(schedule[-1][-1] + 1, episode.length) for i in range(max_idx): input_frames.append(episode.get_frame(i, "video.head_left")) save_mp4( os.path.join(args.output_dir, f"ep{args.episode:06d}_input_head_left.mp4"), np.stack(input_frames, axis=0), fps=args.fps, ) print(f"Avg inference time per chunk: {np.mean(times):.2f}s") print(f"Done. Output dir: {os.path.abspath(args.output_dir)}") def main() -> None: p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) p.add_argument("--model_path", required=True, help="checkpoint-5000 directory") p.add_argument( "--dataset_path", default="/scratch1/home/zhicao/dreamzero/data/trex_small", ) p.add_argument("--episode", type=int, default=0) p.add_argument("--num_chunks", type=int, default=8, help="Number of causal chunks (more -> longer predicted video)") p.add_argument("--prompt", default="perform the task") p.add_argument("--use_dataset_prompt", action="store_true") p.add_argument("--device", default="cuda:0") p.add_argument("--output_dir", default="results_trex_infer") p.add_argument("--fps", type=int, default=5) p.add_argument("--save_input_clip", action="store_true", help="Also save input head_left frames for comparison") run(p.parse_args()) if __name__ == "__main__": main()