#!/usr/bin/env python3 """Offline synthetic-data acceptance. No robot access; no task-quality claim.""" import argparse import json import os import re from pathlib import Path import subprocess import sys import time def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--height", type=int, default=96) parser.add_argument("--width", type=int, default=128) parser.add_argument("--episodes", type=int, default=3) parser.add_argument("--frames-per-episode", type=int, default=24) parser.add_argument("--batch-size", type=int, default=2) parser.add_argument("--steps", type=int, default=10) parser.add_argument("--chunk-size", type=int, default=8) parser.add_argument("--pretrained-backbone", action="store_true") args = parser.parse_args() root = args.output.resolve() root.mkdir(parents=True, exist_ok=False) os.environ["HF_HUB_OFFLINE"] = "1" os.environ["WANDB_MODE"] = "disabled" import numpy as np import torch from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.configs.video import RGBEncoderConfig torch.set_num_threads(4) assert torch.cuda.is_available() features = { "observation.state": {"dtype": "float32", "shape": (6,), "names": [f"joint_{i}" for i in range(6)]}, "action": {"dtype": "float32", "shape": (6,), "names": [f"joint_{i}" for i in range(6)]}, "observation.images.front": {"dtype": "video", "shape": (args.height, args.width, 3), "names": ["height", "width", "channels"]}, } dataset = LeRobotDataset.create( repo_id="local/p0_synthetic", root=root / "dataset", fps=10, features=features, robot_type="synthetic", use_videos=True, rgb_encoder=RGBEncoderConfig(vcodec="h264"), encoder_threads=2, ) for episode in range(args.episodes): for frame in range(args.frames_per_episode): phase = frame / 10 + episode state = np.sin(phase + np.arange(6)).astype(np.float32) image = np.zeros((args.height, args.width, 3), dtype=np.uint8) image[:, :, 1] = 40 + episode * 50 box = max(8, min(args.height, args.width) // 6) x = frame * max(1, (args.width - box) // max(1, args.frames_per_episode - 1)) y = (args.height - box) // 2 image[y:y+box, x:x+box, 0] = 220 dataset.add_frame({"observation.state": state, "action": np.sin(phase + 0.1 + np.arange(6)).astype(np.float32), "observation.images.front": image, "task": "synthetic pipeline check"}) dataset.save_episode() dataset.finalize() from torchcodec.decoders import VideoDecoder videos = list((root / "dataset").rglob("*.mp4")) assert videos for video in videos: decoded = VideoDecoder(str(video)) assert len(decoded) > 0 and decoded[0].shape == (3, args.height, args.width) frame_count = args.episodes * args.frames_per_episode for backend in ("pyav", "torchcodec"): loaded = LeRobotDataset("local/p0_synthetic", root=root / "dataset", video_backend=backend) assert len(loaded) == frame_count boundary_indices = sorted({0, frame_count - 1} | { i for episode in range(1, args.episodes) for i in (episode * args.frames_per_episode - 1, episode * args.frames_per_episode) }) for index in boundary_indices: item = loaded[index] assert item["observation.images.front"].shape == (3, args.height, args.width) assert torch.isfinite(item["action"]).all() train_wrapper = ( "import json,torch; from pathlib import Path; " "from lerobot.scripts.lerobot_train import main; " "torch.cuda.reset_peak_memory_stats(); main(); " f"Path({str(root / 'training_memory.json')!r}).write_text(json.dumps(" "{'allocated_mib':torch.cuda.max_memory_allocated()/2**20," "'reserved_mib':torch.cuda.max_memory_reserved()/2**20}))" ) command = [sys.executable, "-c", train_wrapper, "--policy.type=act", "--policy.device=cuda", "--policy.push_to_hub=false", f"--policy.chunk_size={args.chunk_size}", f"--policy.n_action_steps={args.chunk_size}", "--dataset.repo_id=local/p0_synthetic", f"--dataset.root={root / 'dataset'}", "--dataset.video_backend=torchcodec", f"--batch_size={args.batch_size}", f"--steps={args.steps}", "--num_workers=0", "--seed=42", "--log_freq=1", "--save_freq=10", "--wandb.enable=false", f"--output_dir={root / 'training'}", "--job_name=p0_synthetic"] if not args.pretrained_backbone: command.append("--policy.pretrained_backbone_weights=null") env = dict(os.environ, OMP_NUM_THREADS="4", MKL_NUM_THREADS="4") start = time.perf_counter() with (root / "train.log").open("w") as log: subprocess.run(command, env=env, stdout=log, stderr=subprocess.STDOUT, check=True) elapsed = time.perf_counter() - start checkpoints = sorted((root / "training" / "checkpoints").glob("[0-9]*/pretrained_model")) assert checkpoints, "No checkpoint saved" from lerobot.policies.act.modeling_act import ACTPolicy from lerobot.policies.factory import make_pre_post_processors checkpoint = checkpoints[-1] policy = ACTPolicy.from_pretrained(checkpoint).to("cuda").eval() pre, post = make_pre_post_processors(policy.config, pretrained_path=checkpoint, preprocessor_overrides={"device_processor": {"device": "cuda"}}) observation = {k: v.unsqueeze(0) for k, v in loaded[0].items() if k.startswith("observation.")} timings = [] torch.cuda.reset_peak_memory_stats() with torch.inference_mode(): for i in range(25): policy.reset() torch.cuda.synchronize() begin = time.perf_counter() action = post(policy.select_action(pre(observation))) torch.cuda.synchronize() assert action.shape == (1, 6) and torch.isfinite(action).all() if i >= 5: timings.append((time.perf_counter() - begin) * 1000) report = {"status": "PASS", "scope": "synthetic software pipeline only; not robot/task validation", "frames": frame_count, "episodes": args.episodes, "train_steps": args.steps, "batch_size": args.batch_size, "chunk_size": args.chunk_size, "image_shape": [args.height, args.width, 3], "pretrained_backbone": args.pretrained_backbone, "training_wall_seconds": elapsed, # Upstream resets CUDA peak counters each step. Its per-step log is rounded to 0.01 GiB. "training_peak_allocated_gib_from_step_logs": max(float(v) for v in re.findall( r"mem_gb:([0-9.]+)", (root / "train.log").read_text())), "training_last_step_memory": json.loads((root / "training_memory.json").read_text()), "checkpoint": str(checkpoint), "inference_ms_p50_p95_p99": np.percentile(timings, [50, 95, 99]).tolist(), "inference_peak_allocated_mib": torch.cuda.max_memory_allocated() / 2**20, "torch": torch.__version__, "cuda": torch.version.cuda, "gpu": torch.cuda.get_device_name(0)} (root / "result.json").write_text(json.dumps(report, indent=2) + "\n") print(json.dumps(report, indent=2)) if __name__ == "__main__": main()