146 lines
7.3 KiB
Python
146 lines
7.3 KiB
Python
#!/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()
|