| |
| """Run native LeRobot training/resume and offline inference on real action data. |
| |
| This is a mechanics check, not a robot rollout or task success benchmark. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import importlib.metadata |
| import json |
| import math |
| from pathlib import Path |
| import re |
| import subprocess |
| import sys |
|
|
|
|
| def run(command, directory, name): |
| with (directory / f"{name}.log").open("w") as log: |
| result = subprocess.run(command, stdout=log, stderr=subprocess.STDOUT) |
| record = {"command": command, "exit_code": result.returncode} |
| (directory / f"{name}_command.json").write_text(json.dumps(record, indent=2) + "\n") |
| if result.returncode: |
| raise RuntimeError(f"{name} failed; see {directory / (name + '.log')}") |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--checkpoint", type=Path, required=True) |
| parser.add_argument("--dataset-root", type=Path, required=True) |
| parser.add_argument("--dataset-repo-id", required=True) |
| parser.add_argument("--mode", choices=["binary", "continuous"], required=True) |
| parser.add_argument("--camera-key", default="observation.images.image") |
| parser.add_argument("--episodes", type=int, nargs="+", default=[0, 1, 2]) |
| parser.add_argument("--steps", type=int, default=10) |
| parser.add_argument("--batch-size", type=int, default=4) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--verify-only", action="store_true") |
| args = parser.parse_args() |
| directory = args.output.resolve() |
| directory.mkdir(parents=True, exist_ok=True) |
| info = json.loads((args.dataset_root / "meta/info.json").read_text()) |
| image = info["features"][args.camera_key] |
| shape = image["shape"] |
| names = image.get("names", []) |
| if names == ["height", "width", "channels"]: |
| shape = [shape[2], shape[0], shape[1]] |
| features = {"observation.images.image": {"type": "VISUAL", "shape": shape}} |
| training_dir = directory / "train" |
| first_model = training_dir / f"checkpoints/{args.steps:06d}/pretrained_model" |
| if not args.verify_only: |
| run([ |
| sys.executable, "-m", "lerobot.scripts.lerobot_train", |
| f"--policy.path={args.checkpoint.resolve()}", f"--policy.gripper_mode={args.mode}", |
| "--policy.device=cuda", "--policy.dtype=bfloat16", "--policy.push_to_hub=false", |
| "--policy.optimizer_lr=0.00001", "--policy.input_features=" + json.dumps(features), |
| "--policy.output_features=null", "--rename_map=" + json.dumps({args.camera_key: "observation.images.image"}), |
| f"--dataset.repo_id={args.dataset_repo_id}", f"--dataset.root={args.dataset_root.resolve()}", |
| "--dataset.video_backend=pyav", "--dataset.episodes=" + json.dumps(args.episodes), |
| f"--output_dir={training_dir}", f"--steps={args.steps}", f"--save_freq={args.steps}", |
| f"--batch_size={args.batch_size}", "--num_workers=0", "--log_freq=1", |
| "--env_eval_freq=0", "--wandb.enable=false", "--seed=123", |
| ], directory, "training") |
| run([ |
| sys.executable, "-m", "lerobot.scripts.lerobot_train", "--resume=true", |
| f"--config_path={first_model / 'train_config.json'}", f"--steps={args.steps + 1}", |
| ], directory, "resume") |
|
|
| import torch |
| from safetensors import safe_open |
| import lerobot_policy_dita |
| from lerobot.configs.policies import PreTrainedConfig |
| from lerobot.datasets.factory import resolve_delta_timestamps |
| from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata |
| from lerobot.policies.factory import make_pre_post_processors |
| from lerobot_policy_dita import DitaPolicy |
| from lerobot_policy_dita.policy.processor_dita import DitaMaskedNormalizerProcessorStep, DitaMaskedUnnormalizerProcessorStep |
|
|
| torch.set_num_threads(4) |
| model_dir = training_dir / f"checkpoints/{args.steps + 1:06d}/pretrained_model" |
| cfg = PreTrainedConfig.from_pretrained(model_dir) |
| assert cfg.gripper_mode == args.mode |
| pre, post = make_pre_post_processors(cfg, pretrained_path=model_dir) |
| encoder = next(s for s in pre.steps if isinstance(s, DitaMaskedNormalizerProcessorStep)) |
| decoder = next(s for s in post.steps if isinstance(s, DitaMaskedUnnormalizerProcessorStep)) |
| assert encoder.gripper_mode == decoder.gripper_mode == args.mode |
| raw_stats = json.loads((args.dataset_root / "meta/stats.json").read_text())["action"] |
| |
| |
| initial_pre, initial_post = make_pre_post_processors(cfg, pretrained_path=first_model, |
| preprocessor_overrides={"device_processor": {"device": "cpu"}}) |
| initial_encoder = next(s for s in initial_pre.steps if isinstance(s, DitaMaskedNormalizerProcessorStep)) |
| for key, value in initial_encoder._tensor_stats["action"].items(): |
| torch.testing.assert_close(encoder._tensor_stats["action"][key].cpu().reshape(-1), value.cpu().reshape(-1), rtol=0, atol=0) |
| torch.testing.assert_close(decoder._tensor_stats["action"][key].cpu().reshape(-1), value.cpu().reshape(-1), rtol=0, atol=0) |
| del initial_pre, initial_post, initial_encoder |
| meta = LeRobotDatasetMetadata(args.dataset_repo_id, root=args.dataset_root) |
| dataset = LeRobotDataset(args.dataset_repo_id, root=args.dataset_root, episodes=args.episodes, |
| delta_timestamps=resolve_delta_timestamps(cfg, meta), video_backend="pyav") |
| indices = [0, len(dataset)//4, len(dataset)//2, 3*len(dataset)//4] |
| raw = torch.utils.data.default_collate([dataset[i] for i in indices]) |
| batch = pre(raw) |
| expected = raw["action"][..., -1].to(batch["action"]) |
| if args.mode == "binary": |
| torch.testing.assert_close(batch["action"][..., -1], (1 - expected) / 2) |
| reconstructed = batch["action"].clone() |
| reconstructed[..., -1] = torch.where(reconstructed[..., -1] > cfg.gripper_threshold, 1., -1.) |
| else: |
| lo = encoder._tensor_stats["action"]["min"][-1].to(expected) |
| hi = encoder._tensor_stats["action"]["max"][-1].to(expected) |
| torch.testing.assert_close(batch["action"][..., -1], 2 * (expected - lo) / (hi - lo) - 1) |
| reconstructed = batch["action"] |
| torch.testing.assert_close(post(reconstructed).cpu(), raw["action"], atol=1e-5, rtol=1e-5) |
|
|
| policy = DitaPolicy.from_pretrained(model_dir, config=cfg, strict=True) |
| policy.eval() |
| with torch.no_grad(): |
| loss, _ = policy(batch) |
| assert torch.isfinite(loss) |
| torch.manual_seed(7123) |
| model_actions = policy.predict_action_chunk(batch) |
| actions = post(model_actions) |
| assert actions.shape == (4, cfg.chunk_size, 7) and torch.isfinite(actions).all() |
| gripper = actions[..., -1] |
| if args.mode == "binary": |
| assert set(gripper.flatten().tolist()) <= {-1., 1.} |
| else: |
| assert ((gripper >= lo.cpu() - 1e-5) & (gripper <= hi.cpu() + 1e-5)).all() |
| assert ((gripper > lo.cpu() + 1e-4) & (gripper < hi.cpu() - 1e-4)).any(), "All gripper predictions saturated" |
| |
| policy.reset() |
| online = {k: v for k, v in raw.items() if k != "action" and k != "action_is_pad"} |
| for key in list(online): |
| if key.startswith("observation.images."): |
| online[key] = online[key][:, -1] |
| selected = post(policy.select_action(pre(online))) |
| assert selected.shape == (4, 7) and torch.isfinite(selected).all() |
|
|
| losses, gradients = [], [] |
| for name, count in [("training", args.steps), ("resume", 1)]: |
| log = (directory / f"{name}.log").read_text() |
| found = re.findall(r"step:(\d+).*?loss:(\S+) grdn:(\S+)", log) |
| assert len(found) == count, (name, found) |
| losses.extend(float(x[1]) for x in found) |
| gradients.extend(float(x[2]) for x in found) |
| assert "End of training" in log and "Traceback" not in log |
| assert all(math.isfinite(x) for x in losses + gradients) |
| changed = 0 |
| with safe_open(args.checkpoint / "model.safetensors", framework="pt") as before, safe_open(model_dir / "model.safetensors", framework="pt") as after: |
| for key in after.keys(): |
| value = after.get_tensor(key) |
| assert torch.isfinite(value).all(), key |
| changed += not torch.equal(before.get_tensor(key), value) |
| assert changed > 0 |
| optimizer_steps = set() |
| with safe_open(model_dir.parent / "training_state/optimizer_state.safetensors", framework="pt") as state: |
| for key in state.keys(): |
| value = state.get_tensor(key) |
| assert torch.isfinite(value).all(), key |
| if key.endswith("/step"): |
| optimizer_steps.add(value.item()) |
| assert optimizer_steps == {float(args.steps + 1)}, optimizer_steps |
| report = { |
| "status": "passed", "scope": "Real-data training, resume and offline inference; no task success or hardware rollout claim", |
| "mode": args.mode, "source_checkpoint": str(args.checkpoint.resolve()), |
| "dataset_repo_id": args.dataset_repo_id, "dataset_root": str(args.dataset_root.resolve()), |
| "episodes": args.episodes, "frames": len(dataset), "sample_indices": indices, |
| "training_steps": args.steps, "resume_steps": 1, "batch_size": args.batch_size, |
| "losses": losses, "gradient_norms": gradients, "reloaded_loss": loss.item(), |
| "changed_weight_tensors": changed, "all_model_and_optimizer_tensors_finite": True, |
| "optimizer_steps": sorted(optimizer_steps), "processor_stats_preserved_on_resume": True, |
| "real_action_encode_decode_roundtrip": True, "gripper_bounds_source": "action.min/max" if args.mode == "continuous" else "LIBERO signed commands", |
| "dataset_gripper_min_max": [raw_stats['min'][-1], raw_stats['max'][-1]], |
| "bound_gripper_min_max": [encoder._tensor_stats['action'][k][-1].item() for k in ['min','max']], |
| "raw_gripper_targets": sorted(set(raw['action'][..., -1].flatten().tolist())), |
| "inference_shape": list(actions.shape), "selected_action_shape": list(selected.shape), |
| "inference_gripper_values": gripper.flatten().tolist(), "checkpoint": str(model_dir), |
| "runtime_source": lerobot_policy_dita.__file__, "runtime_version": lerobot_policy_dita.__version__, |
| "versions": {n: importlib.metadata.version(n) for n in ['lerobot','torch','accelerate','diffusers']}, |
| } |
| (directory / "verification.json").write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(report, indent=2), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
