Validate RTX 3080 LeRobot environment
This commit is contained in:
ylhan
Vendored
+145
@@ -0,0 +1,145 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user