File size: 10,723 Bytes
219a93a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | #!/usr/bin/env python3
"""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"]
# Episode subsets can produce different statistics. Compare the initial and
# resumed saved processors; these must agree exactly.
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"
# Also exercise the online single-frame history and select_action path.
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()
|