chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,286 @@
|
||||
"""Train a LeRobot ACT policy using the Rerun dataloader.
|
||||
|
||||
Demonstrates how to stream robot trajectory data from Rerun's catalog
|
||||
into an imitation learning policy (Action Chunking Transformers).
|
||||
|
||||
The Rerun dataloader's Field.window feature fetches future action chunks in a single query per batch.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
|
||||
from lerobot.policies.act.configuration_act import ACTConfig
|
||||
from lerobot.policies.act.modeling_act import ACTPolicy
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from rerun._tracing import tracing_scope, with_tracing
|
||||
from rerun.catalog import CatalogClient
|
||||
from rerun.experimental.dataloader import (
|
||||
DataSource,
|
||||
Field,
|
||||
NumericDecoder,
|
||||
RerunIterableDataset,
|
||||
RerunMapDataset,
|
||||
VideoFrameDecoder,
|
||||
)
|
||||
|
||||
CHECKPOINT_DIR = Path(__file__).resolve().parent / "act_checkpoint"
|
||||
|
||||
IMAGE_H = 32
|
||||
IMAGE_W = 128
|
||||
CAMERAS = ("laptop", "phone", "side")
|
||||
IMAGE_KEYS = tuple(f"observation.images.{cam}" for cam in CAMERAS)
|
||||
|
||||
CHUNK_SIZE = 50
|
||||
EPOCHS = 5
|
||||
BATCH_SIZE = 8
|
||||
LR = 1e-5
|
||||
NUM_WORKERS = 4
|
||||
FETCH_SIZE = 256
|
||||
|
||||
|
||||
class CollateFn:
|
||||
"""Picklable collate callable for PyTorch DataLoader multiprocessing."""
|
||||
|
||||
def __init__(self, chunk_size: int, state_dim: int) -> None:
|
||||
self.chunk_size = chunk_size
|
||||
self.state_dim = state_dim
|
||||
|
||||
@with_tracing("CollateFn")
|
||||
def __call__(self, samples: list[dict[str, torch.Tensor | None]]) -> dict[str, torch.Tensor]:
|
||||
# `VideoFrameDecoder` returns `None` when a target precedes the first keyframe; filter those out.
|
||||
complete: list[dict[str, torch.Tensor]] = [
|
||||
cast("dict[str, torch.Tensor]", s) for s in samples if all(s[f"image_{cam}"] is not None for cam in CAMERAS)
|
||||
]
|
||||
batch_size = len(complete)
|
||||
|
||||
states = torch.stack([s["state"] for s in complete]).float()
|
||||
|
||||
# Future action chunks: reshape windowed flat tensors
|
||||
actions = torch.stack([s["action"].reshape(self.chunk_size, self.state_dim) for s in complete]).float()
|
||||
|
||||
batch: dict[str, torch.Tensor] = {
|
||||
"observation.state": states,
|
||||
"action": actions,
|
||||
"action_is_pad": torch.zeros(batch_size, self.chunk_size, dtype=torch.bool),
|
||||
}
|
||||
# Per-camera images: (3, H, W) uint8 -> float in [0, 1], resized to (IMAGE_H, IMAGE_W)
|
||||
for cam, key in zip(CAMERAS, IMAGE_KEYS):
|
||||
imgs = torch.stack([s[f"image_{cam}"] for s in complete]).float() / 255.0
|
||||
batch[key] = F.interpolate(imgs, size=(IMAGE_H, IMAGE_W), mode="bilinear", align_corners=False)
|
||||
return batch
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument(
|
||||
"--catalog-url",
|
||||
default="rerun+http://127.0.0.1:51234",
|
||||
help="Rerun catalog URL",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset",
|
||||
default="rerun_so101-pick-and-place",
|
||||
help="Dataset name in the catalog",
|
||||
)
|
||||
parser.add_argument("--num-segments", type=int, default=3, help="Number of segments to use (0 for all)")
|
||||
parser.add_argument("--epochs", type=int, default=EPOCHS, help="Number of training epochs")
|
||||
parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="Training batch size")
|
||||
parser.add_argument("--num-workers", type=int, default=NUM_WORKERS, help="DataLoader worker processes")
|
||||
parser.add_argument(
|
||||
"--fetch-size",
|
||||
type=int,
|
||||
default=FETCH_SIZE,
|
||||
help="Samples fetched per server query for the iterable dataset",
|
||||
)
|
||||
parser.add_argument("--lr", type=float, default=LR, help="Learning rate")
|
||||
parser.add_argument(
|
||||
"--dataset-style",
|
||||
choices=("iterable", "map"),
|
||||
default="iterable",
|
||||
help="Which Rerun dataset class to use: 'iterable' (RerunIterableDataset, internal shuffling) "
|
||||
"or 'map' (RerunMapDataset, random access via DataLoader samplers).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--checkpoint-dir",
|
||||
type=Path,
|
||||
default=CHECKPOINT_DIR,
|
||||
help="Directory to save the trained policy checkpoint",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@with_tracing("main")
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
client = CatalogClient(args.catalog_url)
|
||||
dataset_entry = client.get_dataset(args.dataset)
|
||||
|
||||
all_segments = dataset_entry.segment_ids()
|
||||
segments = all_segments if args.num_segments == 0 else all_segments[: args.num_segments]
|
||||
print(f"Using {len(segments)} segments")
|
||||
|
||||
source = DataSource(dataset_entry, segments=segments)
|
||||
|
||||
fields = {
|
||||
"state": Field("/observation.state:Scalars:scalars", decode=NumericDecoder()),
|
||||
"action": Field(
|
||||
"/action:Scalars:scalars",
|
||||
decode=NumericDecoder(),
|
||||
window=(1, CHUNK_SIZE),
|
||||
),
|
||||
"image_laptop": Field(
|
||||
"/observation.images.laptop:VideoStream:sample",
|
||||
decode=VideoFrameDecoder(codec="av1", keyframe_interval=2),
|
||||
),
|
||||
"image_phone": Field(
|
||||
"/observation.images.phone:VideoStream:sample",
|
||||
decode=VideoFrameDecoder(codec="av1", keyframe_interval=2),
|
||||
),
|
||||
"image_side": Field(
|
||||
"/observation.images.side:VideoStream:sample",
|
||||
decode=VideoFrameDecoder(codec="av1", keyframe_interval=2),
|
||||
),
|
||||
}
|
||||
|
||||
ds: RerunIterableDataset | RerunMapDataset
|
||||
if args.dataset_style == "map":
|
||||
ds = RerunMapDataset(source=source, index="frame_index", fields=fields)
|
||||
else:
|
||||
ds = RerunIterableDataset(source=source, index="frame_index", fields=fields, fetch_size=args.fetch_size)
|
||||
print(f"Using {args.dataset_style} dataset with {len(ds)} samples (after window trimming)")
|
||||
|
||||
# IterableDataset doesn't support indexing, so probe shape via iteration.
|
||||
state_tensor = next(iter(ds))["state"]
|
||||
assert state_tensor is not None # NumericDecoder never returns None
|
||||
state_dim = state_tensor.shape[0]
|
||||
action_dim = state_dim
|
||||
print(f"Dimensions: {state_dim=}, {action_dim=}")
|
||||
|
||||
config = ACTConfig(
|
||||
chunk_size=CHUNK_SIZE,
|
||||
n_action_steps=CHUNK_SIZE,
|
||||
use_vae=True,
|
||||
kl_weight=10.0,
|
||||
dim_model=256,
|
||||
n_heads=8,
|
||||
dim_feedforward=1024,
|
||||
n_encoder_layers=4,
|
||||
n_decoder_layers=1,
|
||||
latent_dim=32,
|
||||
n_vae_encoder_layers=4,
|
||||
dropout=0.1,
|
||||
vision_backbone="resnet18",
|
||||
pretrained_backbone_weights=None,
|
||||
normalization_mapping={
|
||||
"STATE": NormalizationMode.MEAN_STD,
|
||||
"VISUAL": NormalizationMode.MEAN_STD,
|
||||
"ACTION": NormalizationMode.MEAN_STD,
|
||||
},
|
||||
input_features={
|
||||
"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(state_dim,)),
|
||||
**{key: PolicyFeature(type=FeatureType.VISUAL, shape=(3, IMAGE_H, IMAGE_W)) for key in IMAGE_KEYS},
|
||||
},
|
||||
output_features={
|
||||
"action": PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,)),
|
||||
},
|
||||
)
|
||||
|
||||
policy = ACTPolicy(config)
|
||||
policy.train()
|
||||
policy.to(device)
|
||||
print(f"ACT policy created ({sum(p.numel() for p in policy.parameters()):,} parameters, device={device})")
|
||||
|
||||
optimizer = torch.optim.AdamW(
|
||||
policy.get_optim_params(),
|
||||
lr=args.lr,
|
||||
weight_decay=1e-4,
|
||||
)
|
||||
|
||||
collate_fn = CollateFn(CHUNK_SIZE, state_dim)
|
||||
# For the map-style dataset, shuffling is driven by the DataLoader's default RandomSampler.
|
||||
# Swap in `sampler=DistributedSampler(ds)` (and call `sampler.set_epoch(epoch)` each epoch)
|
||||
# for multi-node training, or plug in any other PyTorch sampler.
|
||||
loader = DataLoader(
|
||||
ds,
|
||||
batch_size=args.batch_size,
|
||||
shuffle=isinstance(ds, RerunMapDataset),
|
||||
num_workers=args.num_workers,
|
||||
collate_fn=collate_fn,
|
||||
persistent_workers=True,
|
||||
prefetch_factor=8,
|
||||
)
|
||||
|
||||
num_batches = len(loader)
|
||||
print(f"\nTraining for {args.epochs} epochs, {num_batches} batches/epoch, batch_size={args.batch_size}\n")
|
||||
|
||||
for epoch in range(args.epochs):
|
||||
with tracing_scope(f"epoch {epoch}"):
|
||||
if isinstance(ds, RerunIterableDataset):
|
||||
ds.set_epoch(epoch)
|
||||
|
||||
total_loss = 0.0
|
||||
total_l1 = 0.0
|
||||
total_kld = 0.0
|
||||
n = 0
|
||||
|
||||
t_last_print = time.perf_counter()
|
||||
data_sum = 0.0
|
||||
model_sum = 0.0
|
||||
t_data_start = time.perf_counter()
|
||||
for batch in loader:
|
||||
data_time = time.perf_counter() - t_data_start
|
||||
|
||||
t_model_start = time.perf_counter()
|
||||
batch = {k: v.to(device) for k, v in batch.items()}
|
||||
loss, loss_dict = policy.forward(batch)
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
model_time = time.perf_counter() - t_model_start
|
||||
|
||||
total_loss += loss.item()
|
||||
total_l1 += loss_dict["l1_loss"]
|
||||
total_kld += loss_dict.get("kld_loss", 0.0)
|
||||
data_sum += data_time
|
||||
model_sum += model_time
|
||||
n += 1
|
||||
|
||||
if n % 10 == 0 or n == 1:
|
||||
now = time.perf_counter()
|
||||
since_last = now - t_last_print
|
||||
t_last_print = now
|
||||
print(
|
||||
f" epoch {epoch + 1}/{args.epochs} batch {n}/{num_batches}"
|
||||
f" loss={loss.item():.4f}"
|
||||
f" data={data_sum:.1f}s model={model_sum:.1f}s"
|
||||
f" since_last={since_last:.1f}s",
|
||||
flush=True,
|
||||
)
|
||||
data_sum = 0.0
|
||||
model_sum = 0.0
|
||||
|
||||
t_data_start = time.perf_counter()
|
||||
|
||||
avg_loss = total_loss / max(n, 1)
|
||||
avg_l1 = total_l1 / max(n, 1)
|
||||
avg_kld = total_kld / max(n, 1)
|
||||
print(f"Epoch {epoch + 1}/{args.epochs} loss={avg_loss:.4f} l1={avg_l1:.4f} kld={avg_kld:.4f}")
|
||||
|
||||
with tracing_scope("save_pretrained"):
|
||||
policy.save_pretrained(str(args.checkpoint_dir))
|
||||
print(f"\nSaved checkpoint to {args.checkpoint_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user