From 93604011f3a7665bf61c14486d213c2b333fe3c1 Mon Sep 17 00:00:00 2001 From: wilsonchenghy <150957380+wilsonchenghy@users.noreply.github.com> Date: Sat, 27 Jun 2026 01:45:04 -0400 Subject: [PATCH] add-custom-act-policy --- autonomy/il/humanoid_act/__init__.py | 16 + autonomy/il/humanoid_act/checkpoint.py | 46 ++ autonomy/il/humanoid_act/config.py | 62 +++ autonomy/il/humanoid_act/dataset.py | 202 ++++++++ autonomy/il/humanoid_act/detr/__init__.py | 5 + autonomy/il/humanoid_act/detr/build.py | 48 ++ .../il/humanoid_act/detr/models/__init__.py | 9 + .../il/humanoid_act/detr/models/backbone.py | 134 +++++ .../il/humanoid_act/detr/models/detr_vae.py | 160 ++++++ .../detr/models/position_encoding.py | 90 ++++ .../humanoid_act/detr/models/transformer.py | 311 ++++++++++++ .../il/humanoid_act/detr/util/__init__.py | 1 + autonomy/il/humanoid_act/detr/util/misc.py | 468 ++++++++++++++++++ autonomy/il/humanoid_act/eval.py | 108 ++++ autonomy/il/humanoid_act/normalize.py | 108 ++++ autonomy/il/humanoid_act/policy.py | 70 +++ autonomy/il/humanoid_act/smoke.py | 79 +++ autonomy/il/humanoid_act/train.py | 161 ++++++ autonomy/il/pyproject.toml | 4 +- .../so101_vial_task/scripts/lerobot_eval.py | 21 +- 20 files changed, 2099 insertions(+), 4 deletions(-) create mode 100644 autonomy/il/humanoid_act/__init__.py create mode 100644 autonomy/il/humanoid_act/checkpoint.py create mode 100644 autonomy/il/humanoid_act/config.py create mode 100644 autonomy/il/humanoid_act/dataset.py create mode 100644 autonomy/il/humanoid_act/detr/__init__.py create mode 100644 autonomy/il/humanoid_act/detr/build.py create mode 100644 autonomy/il/humanoid_act/detr/models/__init__.py create mode 100644 autonomy/il/humanoid_act/detr/models/backbone.py create mode 100644 autonomy/il/humanoid_act/detr/models/detr_vae.py create mode 100644 autonomy/il/humanoid_act/detr/models/position_encoding.py create mode 100644 autonomy/il/humanoid_act/detr/models/transformer.py create mode 100644 autonomy/il/humanoid_act/detr/util/__init__.py create mode 100644 autonomy/il/humanoid_act/detr/util/misc.py create mode 100644 autonomy/il/humanoid_act/eval.py create mode 100644 autonomy/il/humanoid_act/normalize.py create mode 100644 autonomy/il/humanoid_act/policy.py create mode 100644 autonomy/il/humanoid_act/smoke.py create mode 100644 autonomy/il/humanoid_act/train.py diff --git a/autonomy/il/humanoid_act/__init__.py b/autonomy/il/humanoid_act/__init__.py new file mode 100644 index 00000000..eded9e33 --- /dev/null +++ b/autonomy/il/humanoid_act/__init__.py @@ -0,0 +1,16 @@ +from humanoid_act.config import ACTConfig +from humanoid_act.dataset import ACTBatch, ACTChunkDataset, make_act_dataloaders +from humanoid_act.normalize import NormStats, load_or_compute_stats +from humanoid_act.eval import LocalCustomACTPolicy +from humanoid_act.policy import ACTPolicy + +__all__ = [ + "ACTBatch", + "ACTChunkDataset", + "ACTConfig", + "ACTPolicy", + "LocalCustomACTPolicy", + "NormStats", + "load_or_compute_stats", + "make_act_dataloaders", +] diff --git a/autonomy/il/humanoid_act/checkpoint.py b/autonomy/il/humanoid_act/checkpoint.py new file mode 100644 index 00000000..9ba93984 --- /dev/null +++ b/autonomy/il/humanoid_act/checkpoint.py @@ -0,0 +1,46 @@ +"""Save/load humanoid ACT checkpoints.""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +from humanoid_act.config import ACTConfig +from humanoid_act.normalize import NormStats, save_stats +from humanoid_act.policy import ACTPolicy + + +def save_checkpoint( + path: Path, + *, + policy: ACTPolicy, + config: ACTConfig, + stats: NormStats, + step: int, + val_loss: float | None = None, +) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + payload = { + "step": step, + "val_loss": val_loss, + "policy_state_dict": policy.state_dict(), + "optimizer_state_dict": policy.configure_optimizers().state_dict(), + } + torch.save(payload, path) + config.save(path.parent / "config.json") + save_stats(stats, path.parent / "stats.json") + + +def load_policy(checkpoint_path: Path, device: torch.device) -> tuple[ACTPolicy, ACTConfig, NormStats]: + from humanoid_act.normalize import load_stats + + ckpt_dir = checkpoint_path.parent + config = ACTConfig.load(ckpt_dir / "config.json") + stats = load_stats(ckpt_dir / "stats.json") + policy = ACTPolicy(config) + payload = torch.load(checkpoint_path, map_location=device) + policy.load_state_dict(payload["policy_state_dict"]) + policy.to(device) + policy.eval() + return policy, config, stats diff --git a/autonomy/il/humanoid_act/config.py b/autonomy/il/humanoid_act/config.py new file mode 100644 index 00000000..f4832253 --- /dev/null +++ b/autonomy/il/humanoid_act/config.py @@ -0,0 +1,62 @@ +"""Training / model configuration for humanoid ACT.""" + +from __future__ import annotations + +import json +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + + +@dataclass +class ACTConfig: + repo_id: str + state_dim: int + action_dim: int + camera_names: list[str] + chunk_size: int = 50 + kl_weight: float = 10.0 + lr: float = 1e-5 + lr_backbone: float = 1e-5 + weight_decay: float = 1e-4 + backbone: str = "resnet18" + position_embedding: str = "sine" + hidden_dim: int = 512 + dim_feedforward: int = 3200 + enc_layers: int = 4 + dec_layers: int = 7 + nheads: int = 8 + dropout: float = 0.1 + seed: int = 42 + + @classmethod + def from_dataset_meta( + cls, + repo_id: str, + meta: Any, + *, + chunk_size: int = 50, + **overrides: Any, + ) -> ACTConfig: + state_dim = int(meta.features["observation.state"]["shape"][0]) + action_dim = int(meta.features["action"]["shape"][0]) + camera_names = list(meta.camera_keys) + cfg = cls( + repo_id=repo_id, + state_dim=state_dim, + action_dim=action_dim, + camera_names=camera_names, + chunk_size=chunk_size, + ) + for key, value in overrides.items(): + if hasattr(cfg, key): + setattr(cfg, key, value) + return cfg + + def save(self, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(asdict(self), indent=2)) + + @classmethod + def load(cls, path: Path) -> ACTConfig: + return cls(**json.loads(path.read_text())) diff --git a/autonomy/il/humanoid_act/dataset.py b/autonomy/il/humanoid_act/dataset.py new file mode 100644 index 00000000..fb433866 --- /dev/null +++ b/autonomy/il/humanoid_act/dataset.py @@ -0,0 +1,202 @@ +"""LeRobot dataset → ACT training batches (action chunks + padding mask).""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +from torch.utils.data import DataLoader, Dataset + +from humanoid_act.normalize import NormStats, STATE_KEY, ACTION_KEY + + +@dataclass +class ACTBatch: + images: torch.Tensor # (B, num_cams, C, H, W), float32 in [0, 1] + qpos: torch.Tensor # (B, state_dim) + actions: torch.Tensor # (B, chunk_size, action_dim) + is_pad: torch.Tensor # (B, chunk_size), True where padded + + +def _action_chunk_offsets(fps: float, chunk_size: int, camera_keys: list[str]) -> dict[str, list[float]]: + dt = 1.0 / fps + action_offsets = [i * dt for i in range(chunk_size)] + delta: dict[str, list[float]] = { + STATE_KEY: [0.0], + ACTION_KEY: action_offsets, + } + for cam in camera_keys: + delta[cam] = [0.0] + return delta + + +class ACTChunkDataset(Dataset): + """ + Random-frame ACT dataset backed by LeRobot. + + Each sample: + - images at timestep t (all cameras) + - proprio at t + - action chunk [a_t, a_{t+1}, ...] padded to chunk_size + - is_pad mask for positions after episode end + """ + + def __init__( + self, + repo_id: str, + *, + root: str | None = None, + episodes: list[int] | None = None, + chunk_size: int = 100, + camera_keys: list[str] | None = None, + stats: NormStats | None = None, + ) -> None: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + # Metadata-only load to discover cameras/fps before building deltas. + meta_ds = LeRobotDataset(repo_id, root=root, episodes=episodes) + self.camera_keys = camera_keys or list(meta_ds.meta.camera_keys) + if not self.camera_keys: + raise ValueError("No camera keys found in dataset metadata") + + self.chunk_size = chunk_size + self.stats = stats + self.fps = float(meta_ds.meta.fps) + + offsets = _action_chunk_offsets(self.fps, chunk_size, self.camera_keys) + self._ds = LeRobotDataset( + repo_id, + root=root, + episodes=episodes, + delta_timestamps=offsets, + ) + selected = episodes if episodes is not None else list(range(self._ds.meta.total_episodes)) + self._episode_lengths = { + ep: int(self._ds.meta.episodes[ep]["length"]) for ep in selected + } + + def __len__(self) -> int: + return len(self._ds) + + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: + sample = self._ds[idx] + ep = int(_tensor_scalar(sample["episode_index"])) + frame = int(_tensor_scalar(sample["frame_index"])) + ep_len = self._episode_lengths[ep] + steps_left = max(ep_len - frame, 0) + is_pad = torch.zeros(self.chunk_size, dtype=torch.bool) + if steps_left < self.chunk_size: + is_pad[steps_left:] = True + + images = [] + for cam in self.camera_keys: + img = _to_numpy(sample[cam]).astype(np.float32) + # LeRobot: (T, C, H, W) with T=1 when delta is [0] + if img.ndim == 4: + img = img[0] + images.append(img) + image_stack = np.stack(images, axis=0) / 255.0 # (num_cams, C, H, W) + + qpos = _to_numpy(sample[STATE_KEY]).astype(np.float32).reshape(-1) + if qpos.ndim == 2: + qpos = qpos[0] + + actions = _to_numpy(sample[ACTION_KEY]).astype(np.float32) + + if self.stats is not None: + qpos = self.stats.normalize_state(qpos) + actions = self.stats.normalize_action(actions) + + return { + "images": torch.from_numpy(image_stack), + "qpos": torch.from_numpy(qpos), + "actions": torch.from_numpy(actions), + "is_pad": is_pad, + } + + +def collate_act_batch(items: list[dict[str, torch.Tensor]]) -> ACTBatch: + return ACTBatch( + images=torch.stack([x["images"] for x in items], dim=0), + qpos=torch.stack([x["qpos"] for x in items], dim=0), + actions=torch.stack([x["actions"] for x in items], dim=0), + is_pad=torch.stack([x["is_pad"] for x in items], dim=0), + ) + + +def split_episodes(num_episodes: int, train_ratio: float = 0.8, seed: int = 42) -> tuple[list[int], list[int]]: + rng = np.random.default_rng(seed) + indices = rng.permutation(num_episodes).tolist() + split = int(train_ratio * num_episodes) + return indices[:split], indices[split:] + + +def make_act_dataloaders( + repo_id: str, + *, + root: str | None = None, + chunk_size: int = 100, + batch_size: int = 8, + train_ratio: float = 0.8, + seed: int = 42, + stats: NormStats | None = None, + num_workers: int = 0, +) -> tuple[DataLoader, DataLoader, NormStats]: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + probe = LeRobotDataset(repo_id, root=root) + train_eps, val_eps = split_episodes(probe.num_episodes, train_ratio=train_ratio, seed=seed) + + if stats is None: + from humanoid_act.normalize import load_or_compute_stats + + stats = load_or_compute_stats( + LeRobotDataset(repo_id, root=root, episodes=train_eps), + ) + + train_ds = ACTChunkDataset( + repo_id, + root=root, + episodes=train_eps, + chunk_size=chunk_size, + stats=stats, + ) + val_ds = ACTChunkDataset( + repo_id, + root=root, + episodes=val_eps, + chunk_size=chunk_size, + stats=stats, + ) + + train_loader = DataLoader( + train_ds, + batch_size=batch_size, + shuffle=True, + num_workers=num_workers, + collate_fn=collate_act_batch, + pin_memory=torch.cuda.is_available(), + ) + val_loader = DataLoader( + val_ds, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + collate_fn=collate_act_batch, + pin_memory=torch.cuda.is_available(), + ) + return train_loader, val_loader, stats + + +def _tensor_scalar(value: Any) -> int | float: + if isinstance(value, torch.Tensor): + return value.item() + return value + + +def _to_numpy(value: Any) -> np.ndarray: + if isinstance(value, torch.Tensor): + return value.detach().cpu().numpy() + return np.asarray(value) diff --git a/autonomy/il/humanoid_act/detr/__init__.py b/autonomy/il/humanoid_act/detr/__init__.py new file mode 100644 index 00000000..f4392475 --- /dev/null +++ b/autonomy/il/humanoid_act/detr/__init__.py @@ -0,0 +1,5 @@ +"""Vendored ACT transformer backbone (from tonyzhaozh/act).""" + +from humanoid_act.detr.build import build_act_model_and_optimizer + +__all__ = ["build_act_model_and_optimizer"] diff --git a/autonomy/il/humanoid_act/detr/build.py b/autonomy/il/humanoid_act/detr/build.py new file mode 100644 index 00000000..03bedd95 --- /dev/null +++ b/autonomy/il/humanoid_act/detr/build.py @@ -0,0 +1,48 @@ +"""Build ACT model + optimizer without argparse side effects.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from humanoid_act.detr.models import build_ACT_model + + +def build_act_model_and_optimizer(config) -> tuple[torch.nn.Module, torch.optim.Optimizer]: + args = SimpleNamespace( + lr=config.lr, + lr_backbone=config.lr_backbone, + weight_decay=config.weight_decay, + backbone=config.backbone, + dilation=False, + position_embedding=config.position_embedding, + camera_names=list(config.camera_names), + enc_layers=config.enc_layers, + dec_layers=config.dec_layers, + dim_feedforward=config.dim_feedforward, + hidden_dim=config.hidden_dim, + dropout=config.dropout, + nheads=config.nheads, + num_queries=config.chunk_size, + pre_norm=False, + masks=False, + state_dim=config.state_dim, + ) + + model = build_ACT_model(args) + param_dicts = [ + { + "params": [ + p for n, p in model.named_parameters() if "backbone" not in n and p.requires_grad + ] + }, + { + "params": [ + p for n, p in model.named_parameters() if "backbone" in n and p.requires_grad + ], + "lr": args.lr_backbone, + }, + ] + optimizer = torch.optim.AdamW(param_dicts, lr=args.lr, weight_decay=args.weight_decay) + return model, optimizer diff --git a/autonomy/il/humanoid_act/detr/models/__init__.py b/autonomy/il/humanoid_act/detr/models/__init__.py new file mode 100644 index 00000000..8bdb8c69 --- /dev/null +++ b/autonomy/il/humanoid_act/detr/models/__init__.py @@ -0,0 +1,9 @@ +from .detr_vae import build, build_cnnmlp + + +def build_ACT_model(args): + return build(args) + + +def build_CNNMLP_model(args): + return build_cnnmlp(args) diff --git a/autonomy/il/humanoid_act/detr/models/backbone.py b/autonomy/il/humanoid_act/detr/models/backbone.py new file mode 100644 index 00000000..618bb9ec --- /dev/null +++ b/autonomy/il/humanoid_act/detr/models/backbone.py @@ -0,0 +1,134 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Backbone modules. +""" +from collections import OrderedDict + +import torch +import torch.nn.functional as F +import torchvision +from torch import nn +from torchvision.models._utils import IntermediateLayerGetter +from typing import Dict, List + +from ..util.misc import NestedTensor, is_main_process + +from .position_encoding import build_position_encoding + +class FrozenBatchNorm2d(torch.nn.Module): + """ + BatchNorm2d where the batch statistics and the affine parameters are fixed. + + Copy-paste from torchvision.misc.ops with added eps before rqsrt, + without which any other policy_models than torchvision.policy_models.resnet[18,34,50,101] + produce nans. + """ + + def __init__(self, n): + super(FrozenBatchNorm2d, self).__init__() + self.register_buffer("weight", torch.ones(n)) + self.register_buffer("bias", torch.zeros(n)) + self.register_buffer("running_mean", torch.zeros(n)) + self.register_buffer("running_var", torch.ones(n)) + + def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, + missing_keys, unexpected_keys, error_msgs): + num_batches_tracked_key = prefix + 'num_batches_tracked' + if num_batches_tracked_key in state_dict: + del state_dict[num_batches_tracked_key] + + super(FrozenBatchNorm2d, self)._load_from_state_dict( + state_dict, prefix, local_metadata, strict, + missing_keys, unexpected_keys, error_msgs) + + def forward(self, x): + # move reshapes to the beginning + # to make it fuser-friendly + w = self.weight.reshape(1, -1, 1, 1) + b = self.bias.reshape(1, -1, 1, 1) + rv = self.running_var.reshape(1, -1, 1, 1) + rm = self.running_mean.reshape(1, -1, 1, 1) + eps = 1e-5 + scale = w * (rv + eps).rsqrt() + bias = b - rm * scale + return x * scale + bias + + +class BackboneBase(nn.Module): + + def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool): + super().__init__() + # for name, parameter in backbone.named_parameters(): # only train later layers # TODO do we want this? + # if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name: + # parameter.requires_grad_(False) + if return_interm_layers: + return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"} + else: + return_layers = {'layer4': "0"} + self.body = IntermediateLayerGetter(backbone, return_layers=return_layers) + self.num_channels = num_channels + + def forward(self, tensor): + xs = self.body(tensor) + return xs + # out: Dict[str, NestedTensor] = {} + # for name, x in xs.items(): + # m = tensor_list.mask + # assert m is not None + # mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0] + # out[name] = NestedTensor(x, mask) + # return out + + +class Backbone(BackboneBase): + """ResNet backbone with frozen BatchNorm.""" + def __init__(self, name: str, + train_backbone: bool, + return_interm_layers: bool, + dilation: bool): + factory = getattr(torchvision.models, name) + kwargs = { + "replace_stride_with_dilation": [False, False, dilation], + "norm_layer": FrozenBatchNorm2d, + } + if is_main_process(): + try: + from torchvision.models import ResNet18_Weights, ResNet34_Weights, ResNet50_Weights + + weights_map = { + "resnet18": ResNet18_Weights.DEFAULT, + "resnet34": ResNet34_Weights.DEFAULT, + "resnet50": ResNet50_Weights.DEFAULT, + } + kwargs["weights"] = weights_map[name] + except ImportError: + kwargs["pretrained"] = True + backbone = factory(**kwargs) + num_channels = 512 if name in ('resnet18', 'resnet34') else 2048 + super().__init__(backbone, train_backbone, num_channels, return_interm_layers) + + +class Joiner(nn.Sequential): + def __init__(self, backbone, position_embedding): + super().__init__(backbone, position_embedding) + + def forward(self, tensor_list: NestedTensor): + xs = self[0](tensor_list) + out: List[NestedTensor] = [] + pos = [] + for name, x in xs.items(): + out.append(x) + # position encoding + pos.append(self[1](x).to(x.dtype)) + + return out, pos + + +def build_backbone(args): + position_embedding = build_position_encoding(args) + train_backbone = args.lr_backbone > 0 + return_interm_layers = args.masks + backbone = Backbone(args.backbone, train_backbone, return_interm_layers, args.dilation) + model = Joiner(backbone, position_embedding) + model.num_channels = backbone.num_channels + return model diff --git a/autonomy/il/humanoid_act/detr/models/detr_vae.py b/autonomy/il/humanoid_act/detr/models/detr_vae.py new file mode 100644 index 00000000..390081ee --- /dev/null +++ b/autonomy/il/humanoid_act/detr/models/detr_vae.py @@ -0,0 +1,160 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +"""DETR-VAE model for ACT (adapted for configurable state/action dims).""" + +from __future__ import annotations + +import numpy as np +import torch +from torch import nn +from torch.autograd import Variable + +from .backbone import build_backbone +from .transformer import TransformerEncoder, TransformerEncoderLayer, build_transformer + + +def reparametrize(mu, logvar): + std = logvar.div(2).exp() + eps = Variable(std.data.new(std.size()).normal_()) + return mu + std * eps + + +def get_sinusoid_encoding_table(n_position, d_hid): + def get_position_angle_vec(position): + return [ + position / np.power(10000, 2 * (hid_j // 2) / d_hid) + for hid_j in range(d_hid) + ] + + sinusoid_table = np.array( + [get_position_angle_vec(pos_i) for pos_i in range(n_position)] + ) + sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) + sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) + return torch.FloatTensor(sinusoid_table).unsqueeze(0) + + +class DETRVAE(nn.Module): + def __init__( + self, + backbones, + transformer, + encoder, + *, + state_dim: int, + num_queries: int, + camera_names: list[str], + ): + super().__init__() + self.num_queries = num_queries + self.camera_names = camera_names + self.state_dim = state_dim + self.transformer = transformer + self.encoder = encoder + hidden_dim = transformer.d_model + self.action_head = nn.Linear(hidden_dim, state_dim) + self.is_pad_head = nn.Linear(hidden_dim, 1) + self.query_embed = nn.Embedding(num_queries, hidden_dim) + self.input_proj = nn.Conv2d(backbones[0].num_channels, hidden_dim, kernel_size=1) + self.backbones = nn.ModuleList(backbones) + self.input_proj_robot_state = nn.Linear(state_dim, hidden_dim) + + self.latent_dim = 32 + self.cls_embed = nn.Embedding(1, hidden_dim) + self.encoder_action_proj = nn.Linear(state_dim, hidden_dim) + self.encoder_joint_proj = nn.Linear(state_dim, hidden_dim) + self.latent_proj = nn.Linear(hidden_dim, self.latent_dim * 2) + self.register_buffer( + "pos_table", + get_sinusoid_encoding_table(1 + 1 + num_queries, hidden_dim), + ) + self.latent_out_proj = nn.Linear(self.latent_dim, hidden_dim) + self.additional_pos_embed = nn.Embedding(2, hidden_dim) + + def forward(self, qpos, image, env_state, actions=None, is_pad=None): + del env_state + is_training = actions is not None + bs, _ = qpos.shape + + if is_training: + action_embed = self.encoder_action_proj(actions) + qpos_embed = self.encoder_joint_proj(qpos).unsqueeze(1) + cls_embed = self.cls_embed.weight.unsqueeze(0).repeat(bs, 1, 1) + encoder_input = torch.cat([cls_embed, qpos_embed, action_embed], axis=1) + encoder_input = encoder_input.permute(1, 0, 2) + cls_joint_is_pad = torch.full((bs, 2), False, device=qpos.device) + is_pad = torch.cat([cls_joint_is_pad, is_pad], axis=1) + pos_embed = self.pos_table.clone().detach().permute(1, 0, 2) + encoder_output = self.encoder( + encoder_input, pos=pos_embed, src_key_padding_mask=is_pad + )[0] + latent_info = self.latent_proj(encoder_output) + mu = latent_info[:, : self.latent_dim] + logvar = latent_info[:, self.latent_dim :] + latent_sample = reparametrize(mu, logvar) + latent_input = self.latent_out_proj(latent_sample) + else: + mu = logvar = None + latent_sample = torch.zeros([bs, self.latent_dim], dtype=torch.float32, device=qpos.device) + latent_input = self.latent_out_proj(latent_sample) + + all_cam_features = [] + all_cam_pos = [] + for cam_id, _cam_name in enumerate(self.camera_names): + features, pos = self.backbones[0](image[:, cam_id]) + features = features[0] + pos = pos[0] + all_cam_features.append(self.input_proj(features)) + all_cam_pos.append(pos) + + proprio_input = self.input_proj_robot_state(qpos) + src = torch.cat(all_cam_features, axis=3) + pos = torch.cat(all_cam_pos, axis=3) + hs = self.transformer( + src, + None, + self.query_embed.weight, + pos, + latent_input, + proprio_input, + self.additional_pos_embed.weight, + )[0] + a_hat = self.action_head(hs) + is_pad_hat = self.is_pad_head(hs) + return a_hat, is_pad_hat, [mu, logvar] + + +def build_encoder(args): + d_model = args.hidden_dim + dropout = args.dropout + nhead = args.nheads + dim_feedforward = args.dim_feedforward + num_encoder_layers = args.enc_layers + normalize_before = args.pre_norm + activation = "relu" + + encoder_layer = TransformerEncoderLayer( + d_model, nhead, dim_feedforward, dropout, activation, normalize_before + ) + encoder_norm = nn.LayerNorm(d_model) if normalize_before else None + return TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) + + +def build(args): + backbones = [build_backbone(args)] + transformer = build_transformer(args) + encoder = build_encoder(args) + model = DETRVAE( + backbones, + transformer, + encoder, + state_dim=args.state_dim, + num_queries=args.num_queries, + camera_names=args.camera_names, + ) + n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f"number of parameters: {n_parameters / 1e6:.2f}M") + return model + + +def build_cnnmlp(args): + raise NotImplementedError("CNNMLP is not wired up in humanoid_act yet") diff --git a/autonomy/il/humanoid_act/detr/models/position_encoding.py b/autonomy/il/humanoid_act/detr/models/position_encoding.py new file mode 100644 index 00000000..db455d9c --- /dev/null +++ b/autonomy/il/humanoid_act/detr/models/position_encoding.py @@ -0,0 +1,90 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Various positional encodings for the transformer. +""" +import math +import torch +from torch import nn + +from ..util.misc import NestedTensor + +class PositionEmbeddingSine(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one + used by the Attention is all you need paper, generalized to work on images. + """ + def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + + def forward(self, tensor): + x = tensor + # mask = tensor_list.mask + # assert mask is not None + # not_mask = ~mask + + not_mask = torch.ones_like(x[0, [0]]) + y_embed = not_mask.cumsum(1, dtype=torch.float32) + x_embed = not_mask.cumsum(2, dtype=torch.float32) + if self.normalize: + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) + pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + return pos + + +class PositionEmbeddingLearned(nn.Module): + """ + Absolute pos embedding, learned. + """ + def __init__(self, num_pos_feats=256): + super().__init__() + self.row_embed = nn.Embedding(50, num_pos_feats) + self.col_embed = nn.Embedding(50, num_pos_feats) + self.reset_parameters() + + def reset_parameters(self): + nn.init.uniform_(self.row_embed.weight) + nn.init.uniform_(self.col_embed.weight) + + def forward(self, tensor_list: NestedTensor): + x = tensor_list.tensors + h, w = x.shape[-2:] + i = torch.arange(w, device=x.device) + j = torch.arange(h, device=x.device) + x_emb = self.col_embed(i) + y_emb = self.row_embed(j) + pos = torch.cat([ + x_emb.unsqueeze(0).repeat(h, 1, 1), + y_emb.unsqueeze(1).repeat(1, w, 1), + ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1) + return pos + + +def build_position_encoding(args): + N_steps = args.hidden_dim // 2 + if args.position_embedding in ('v2', 'sine'): + # TODO find a better way of exposing other arguments + position_embedding = PositionEmbeddingSine(N_steps, normalize=True) + elif args.position_embedding in ('v3', 'learned'): + position_embedding = PositionEmbeddingLearned(N_steps) + else: + raise ValueError(f"not supported {args.position_embedding}") + + return position_embedding diff --git a/autonomy/il/humanoid_act/detr/models/transformer.py b/autonomy/il/humanoid_act/detr/models/transformer.py new file mode 100644 index 00000000..f14278e6 --- /dev/null +++ b/autonomy/il/humanoid_act/detr/models/transformer.py @@ -0,0 +1,311 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +DETR Transformer class. + +Copy-paste from torch.nn.Transformer with modifications: + * positional encodings are passed in MHattention + * extra LN at the end of encoder is removed + * decoder returns a stack of activations from all decoding layers +""" +import copy +from typing import Optional, List + +import torch +import torch.nn.functional as F +from torch import nn, Tensor + +class Transformer(nn.Module): + + def __init__(self, d_model=512, nhead=8, num_encoder_layers=6, + num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, + activation="relu", normalize_before=False, + return_intermediate_dec=False): + super().__init__() + + encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, + dropout, activation, normalize_before) + encoder_norm = nn.LayerNorm(d_model) if normalize_before else None + self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) + + decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, + dropout, activation, normalize_before) + decoder_norm = nn.LayerNorm(d_model) + self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm, + return_intermediate=return_intermediate_dec) + + self._reset_parameters() + + self.d_model = d_model + self.nhead = nhead + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def forward(self, src, mask, query_embed, pos_embed, latent_input=None, proprio_input=None, additional_pos_embed=None): + # TODO flatten only when input has H and W + if len(src.shape) == 4: # has H and W + # flatten NxCxHxW to HWxNxC + bs, c, h, w = src.shape + src = src.flatten(2).permute(2, 0, 1) + pos_embed = pos_embed.flatten(2).permute(2, 0, 1).repeat(1, bs, 1) + query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1) + # mask = mask.flatten(1) + + additional_pos_embed = additional_pos_embed.unsqueeze(1).repeat(1, bs, 1) # seq, bs, dim + pos_embed = torch.cat([additional_pos_embed, pos_embed], axis=0) + + addition_input = torch.stack([latent_input, proprio_input], axis=0) + src = torch.cat([addition_input, src], axis=0) + else: + assert len(src.shape) == 3 + # flatten NxHWxC to HWxNxC + bs, hw, c = src.shape + src = src.permute(1, 0, 2) + pos_embed = pos_embed.unsqueeze(1).repeat(1, bs, 1) + query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1) + + tgt = torch.zeros_like(query_embed) + memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed) + hs = self.decoder(tgt, memory, memory_key_padding_mask=mask, + pos=pos_embed, query_pos=query_embed) + hs = hs.transpose(1, 2) + return hs + +class TransformerEncoder(nn.Module): + + def __init__(self, encoder_layer, num_layers, norm=None): + super().__init__() + self.layers = _get_clones(encoder_layer, num_layers) + self.num_layers = num_layers + self.norm = norm + + def forward(self, src, + mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None): + output = src + + for layer in self.layers: + output = layer(output, src_mask=mask, + src_key_padding_mask=src_key_padding_mask, pos=pos) + + if self.norm is not None: + output = self.norm(output) + + return output + + +class TransformerDecoder(nn.Module): + + def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False): + super().__init__() + self.layers = _get_clones(decoder_layer, num_layers) + self.num_layers = num_layers + self.norm = norm + self.return_intermediate = return_intermediate + + def forward(self, tgt, memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + output = tgt + + intermediate = [] + + for layer in self.layers: + output = layer(output, memory, tgt_mask=tgt_mask, + memory_mask=memory_mask, + tgt_key_padding_mask=tgt_key_padding_mask, + memory_key_padding_mask=memory_key_padding_mask, + pos=pos, query_pos=query_pos) + if self.return_intermediate: + intermediate.append(self.norm(output)) + + if self.norm is not None: + output = self.norm(output) + if self.return_intermediate: + intermediate.pop() + intermediate.append(output) + + if self.return_intermediate: + return torch.stack(intermediate) + + return output.unsqueeze(0) + + +class TransformerEncoderLayer(nn.Module): + + def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, + activation="relu", normalize_before=False): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post(self, + src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None): + q = k = self.with_pos_embed(src, pos) + src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, + key_padding_mask=src_key_padding_mask)[0] + src = src + self.dropout1(src2) + src = self.norm1(src) + src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) + src = src + self.dropout2(src2) + src = self.norm2(src) + return src + + def forward_pre(self, src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None): + src2 = self.norm1(src) + q = k = self.with_pos_embed(src2, pos) + src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask, + key_padding_mask=src_key_padding_mask)[0] + src = src + self.dropout1(src2) + src2 = self.norm2(src) + src2 = self.linear2(self.dropout(self.activation(self.linear1(src2)))) + src = src + self.dropout2(src2) + return src + + def forward(self, src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None): + if self.normalize_before: + return self.forward_pre(src, src_mask, src_key_padding_mask, pos) + return self.forward_post(src, src_mask, src_key_padding_mask, pos) + + +class TransformerDecoderLayer(nn.Module): + + def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, + activation="relu", normalize_before=False): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.norm3 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post(self, tgt, memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + q = k = self.with_pos_embed(tgt, query_pos) + tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask)[0] + tgt = tgt + self.dropout1(tgt2) + tgt = self.norm1(tgt) + tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask)[0] + tgt = tgt + self.dropout2(tgt2) + tgt = self.norm2(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout3(tgt2) + tgt = self.norm3(tgt) + return tgt + + def forward_pre(self, tgt, memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + tgt2 = self.norm1(tgt) + q = k = self.with_pos_embed(tgt2, query_pos) + tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask)[0] + tgt = tgt + self.dropout1(tgt2) + tgt2 = self.norm2(tgt) + tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask)[0] + tgt = tgt + self.dropout2(tgt2) + tgt2 = self.norm3(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) + tgt = tgt + self.dropout3(tgt2) + return tgt + + def forward(self, tgt, memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + if self.normalize_before: + return self.forward_pre(tgt, memory, tgt_mask, memory_mask, + tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) + return self.forward_post(tgt, memory, tgt_mask, memory_mask, + tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) + + +def _get_clones(module, N): + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) + + +def build_transformer(args): + return Transformer( + d_model=args.hidden_dim, + dropout=args.dropout, + nhead=args.nheads, + dim_feedforward=args.dim_feedforward, + num_encoder_layers=args.enc_layers, + num_decoder_layers=args.dec_layers, + normalize_before=args.pre_norm, + return_intermediate_dec=True, + ) + + +def _get_activation_fn(activation): + """Return an activation function given a string""" + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + if activation == "glu": + return F.glu + raise RuntimeError(F"activation should be relu/gelu, not {activation}.") diff --git a/autonomy/il/humanoid_act/detr/util/__init__.py b/autonomy/il/humanoid_act/detr/util/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/autonomy/il/humanoid_act/detr/util/__init__.py @@ -0,0 +1 @@ + diff --git a/autonomy/il/humanoid_act/detr/util/misc.py b/autonomy/il/humanoid_act/detr/util/misc.py new file mode 100644 index 00000000..dfa9fb5b --- /dev/null +++ b/autonomy/il/humanoid_act/detr/util/misc.py @@ -0,0 +1,468 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Misc functions, including distributed helpers. + +Mostly copy-paste from torchvision references. +""" +import os +import subprocess +import time +from collections import defaultdict, deque +import datetime +import pickle +from packaging import version +from typing import Optional, List + +import torch +import torch.distributed as dist +from torch import Tensor + +# needed due to empty tensor bug in pytorch and torchvision 0.5 +import torchvision +if version.parse(torchvision.__version__) < version.parse('0.7'): + from torchvision.ops import _new_empty_tensor + from torchvision.ops.misc import _output_size + + +class SmoothedValue(object): + """Track a series of values and provide access to smoothed values over a + window or the global series average. + """ + + def __init__(self, window_size=20, fmt=None): + if fmt is None: + fmt = "{median:.4f} ({global_avg:.4f})" + self.deque = deque(maxlen=window_size) + self.total = 0.0 + self.count = 0 + self.fmt = fmt + + def update(self, value, n=1): + self.deque.append(value) + self.count += n + self.total += value * n + + def synchronize_between_processes(self): + """ + Warning: does not synchronize the deque! + """ + if not is_dist_avail_and_initialized(): + return + t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda') + dist.barrier() + dist.all_reduce(t) + t = t.tolist() + self.count = int(t[0]) + self.total = t[1] + + @property + def median(self): + d = torch.tensor(list(self.deque)) + return d.median().item() + + @property + def avg(self): + d = torch.tensor(list(self.deque), dtype=torch.float32) + return d.mean().item() + + @property + def global_avg(self): + return self.total / self.count + + @property + def max(self): + return max(self.deque) + + @property + def value(self): + return self.deque[-1] + + def __str__(self): + return self.fmt.format( + median=self.median, + avg=self.avg, + global_avg=self.global_avg, + max=self.max, + value=self.value) + + +def all_gather(data): + """ + Run all_gather on arbitrary picklable data (not necessarily tensors) + Args: + data: any picklable object + Returns: + list[data]: list of data gathered from each rank + """ + world_size = get_world_size() + if world_size == 1: + return [data] + + # serialized to a Tensor + buffer = pickle.dumps(data) + storage = torch.ByteStorage.from_buffer(buffer) + tensor = torch.ByteTensor(storage).to("cuda") + + # obtain Tensor size of each rank + local_size = torch.tensor([tensor.numel()], device="cuda") + size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)] + dist.all_gather(size_list, local_size) + size_list = [int(size.item()) for size in size_list] + max_size = max(size_list) + + # receiving Tensor from all ranks + # we pad the tensor because torch all_gather does not support + # gathering tensors of different shapes + tensor_list = [] + for _ in size_list: + tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda")) + if local_size != max_size: + padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda") + tensor = torch.cat((tensor, padding), dim=0) + dist.all_gather(tensor_list, tensor) + + data_list = [] + for size, tensor in zip(size_list, tensor_list): + buffer = tensor.cpu().numpy().tobytes()[:size] + data_list.append(pickle.loads(buffer)) + + return data_list + + +def reduce_dict(input_dict, average=True): + """ + Args: + input_dict (dict): all the values will be reduced + average (bool): whether to do average or sum + Reduce the values in the dictionary from all processes so that all processes + have the averaged results. Returns a dict with the same fields as + input_dict, after reduction. + """ + world_size = get_world_size() + if world_size < 2: + return input_dict + with torch.no_grad(): + names = [] + values = [] + # sort the keys so that they are consistent across processes + for k in sorted(input_dict.keys()): + names.append(k) + values.append(input_dict[k]) + values = torch.stack(values, dim=0) + dist.all_reduce(values) + if average: + values /= world_size + reduced_dict = {k: v for k, v in zip(names, values)} + return reduced_dict + + +class MetricLogger(object): + def __init__(self, delimiter="\t"): + self.meters = defaultdict(SmoothedValue) + self.delimiter = delimiter + + def update(self, **kwargs): + for k, v in kwargs.items(): + if isinstance(v, torch.Tensor): + v = v.item() + assert isinstance(v, (float, int)) + self.meters[k].update(v) + + def __getattr__(self, attr): + if attr in self.meters: + return self.meters[attr] + if attr in self.__dict__: + return self.__dict__[attr] + raise AttributeError("'{}' object has no attribute '{}'".format( + type(self).__name__, attr)) + + def __str__(self): + loss_str = [] + for name, meter in self.meters.items(): + loss_str.append( + "{}: {}".format(name, str(meter)) + ) + return self.delimiter.join(loss_str) + + def synchronize_between_processes(self): + for meter in self.meters.values(): + meter.synchronize_between_processes() + + def add_meter(self, name, meter): + self.meters[name] = meter + + def log_every(self, iterable, print_freq, header=None): + i = 0 + if not header: + header = '' + start_time = time.time() + end = time.time() + iter_time = SmoothedValue(fmt='{avg:.4f}') + data_time = SmoothedValue(fmt='{avg:.4f}') + space_fmt = ':' + str(len(str(len(iterable)))) + 'd' + if torch.cuda.is_available(): + log_msg = self.delimiter.join([ + header, + '[{0' + space_fmt + '}/{1}]', + 'eta: {eta}', + '{meters}', + 'time: {time}', + 'data: {data}', + 'max mem: {memory:.0f}' + ]) + else: + log_msg = self.delimiter.join([ + header, + '[{0' + space_fmt + '}/{1}]', + 'eta: {eta}', + '{meters}', + 'time: {time}', + 'data: {data}' + ]) + MB = 1024.0 * 1024.0 + for obj in iterable: + data_time.update(time.time() - end) + yield obj + iter_time.update(time.time() - end) + if i % print_freq == 0 or i == len(iterable) - 1: + eta_seconds = iter_time.global_avg * (len(iterable) - i) + eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) + if torch.cuda.is_available(): + print(log_msg.format( + i, len(iterable), eta=eta_string, + meters=str(self), + time=str(iter_time), data=str(data_time), + memory=torch.cuda.max_memory_allocated() / MB)) + else: + print(log_msg.format( + i, len(iterable), eta=eta_string, + meters=str(self), + time=str(iter_time), data=str(data_time))) + i += 1 + end = time.time() + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print('{} Total time: {} ({:.4f} s / it)'.format( + header, total_time_str, total_time / len(iterable))) + + +def get_sha(): + cwd = os.path.dirname(os.path.abspath(__file__)) + + def _run(command): + return subprocess.check_output(command, cwd=cwd).decode('ascii').strip() + sha = 'N/A' + diff = "clean" + branch = 'N/A' + try: + sha = _run(['git', 'rev-parse', 'HEAD']) + subprocess.check_output(['git', 'diff'], cwd=cwd) + diff = _run(['git', 'diff-index', 'HEAD']) + diff = "has uncommited changes" if diff else "clean" + branch = _run(['git', 'rev-parse', '--abbrev-ref', 'HEAD']) + except Exception: + pass + message = f"sha: {sha}, status: {diff}, branch: {branch}" + return message + + +def collate_fn(batch): + batch = list(zip(*batch)) + batch[0] = nested_tensor_from_tensor_list(batch[0]) + return tuple(batch) + + +def _max_by_axis(the_list): + # type: (List[List[int]]) -> List[int] + maxes = the_list[0] + for sublist in the_list[1:]: + for index, item in enumerate(sublist): + maxes[index] = max(maxes[index], item) + return maxes + + +class NestedTensor(object): + def __init__(self, tensors, mask: Optional[Tensor]): + self.tensors = tensors + self.mask = mask + + def to(self, device): + # type: (Device) -> NestedTensor # noqa + cast_tensor = self.tensors.to(device) + mask = self.mask + if mask is not None: + assert mask is not None + cast_mask = mask.to(device) + else: + cast_mask = None + return NestedTensor(cast_tensor, cast_mask) + + def decompose(self): + return self.tensors, self.mask + + def __repr__(self): + return str(self.tensors) + + +def nested_tensor_from_tensor_list(tensor_list: List[Tensor]): + # TODO make this more general + if tensor_list[0].ndim == 3: + if torchvision._is_tracing(): + # nested_tensor_from_tensor_list() does not export well to ONNX + # call _onnx_nested_tensor_from_tensor_list() instead + return _onnx_nested_tensor_from_tensor_list(tensor_list) + + # TODO make it support different-sized images + max_size = _max_by_axis([list(img.shape) for img in tensor_list]) + # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list])) + batch_shape = [len(tensor_list)] + max_size + b, c, h, w = batch_shape + dtype = tensor_list[0].dtype + device = tensor_list[0].device + tensor = torch.zeros(batch_shape, dtype=dtype, device=device) + mask = torch.ones((b, h, w), dtype=torch.bool, device=device) + for img, pad_img, m in zip(tensor_list, tensor, mask): + pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) + m[: img.shape[1], :img.shape[2]] = False + else: + raise ValueError('not supported') + return NestedTensor(tensor, mask) + + +# _onnx_nested_tensor_from_tensor_list() is an implementation of +# nested_tensor_from_tensor_list() that is supported by ONNX tracing. +@torch.jit.unused +def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor: + max_size = [] + for i in range(tensor_list[0].dim()): + max_size_i = torch.max(torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)).to(torch.int64) + max_size.append(max_size_i) + max_size = tuple(max_size) + + # work around for + # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) + # m[: img.shape[1], :img.shape[2]] = False + # which is not yet supported in onnx + padded_imgs = [] + padded_masks = [] + for img in tensor_list: + padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))] + padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0])) + padded_imgs.append(padded_img) + + m = torch.zeros_like(img[0], dtype=torch.int, device=img.device) + padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1) + padded_masks.append(padded_mask.to(torch.bool)) + + tensor = torch.stack(padded_imgs) + mask = torch.stack(padded_masks) + + return NestedTensor(tensor, mask=mask) + + +def setup_for_distributed(is_master): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + builtin_print = __builtin__.print + + def print(*args, **kwargs): + force = kwargs.pop('force', False) + if is_master or force: + builtin_print(*args, **kwargs) + + __builtin__.print = print + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_world_size(): + if not is_dist_avail_and_initialized(): + return 1 + return dist.get_world_size() + + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() + + +def is_main_process(): + return get_rank() == 0 + + +def save_on_master(*args, **kwargs): + if is_main_process(): + torch.save(*args, **kwargs) + + +def init_distributed_mode(args): + if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: + args.rank = int(os.environ["RANK"]) + args.world_size = int(os.environ['WORLD_SIZE']) + args.gpu = int(os.environ['LOCAL_RANK']) + elif 'SLURM_PROCID' in os.environ: + args.rank = int(os.environ['SLURM_PROCID']) + args.gpu = args.rank % torch.cuda.device_count() + else: + print('Not using distributed mode') + args.distributed = False + return + + args.distributed = True + + torch.cuda.set_device(args.gpu) + args.dist_backend = 'nccl' + print('| distributed init (rank {}): {}'.format( + args.rank, args.dist_url), flush=True) + torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url, + world_size=args.world_size, rank=args.rank) + torch.distributed.barrier() + setup_for_distributed(args.rank == 0) + + +@torch.no_grad() +def accuracy(output, target, topk=(1,)): + """Computes the precision@k for the specified values of k""" + if target.numel() == 0: + return [torch.zeros([], device=output.device)] + maxk = max(topk) + batch_size = target.size(0) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + + res = [] + for k in topk: + correct_k = correct[:k].view(-1).float().sum(0) + res.append(correct_k.mul_(100.0 / batch_size)) + return res + + +def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None): + # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor + """ + Equivalent to nn.functional.interpolate, but with support for empty batch sizes. + This will eventually be supported natively by PyTorch, and this + class can go away. + """ + if version.parse(torchvision.__version__) < version.parse('0.7'): + if input.numel() > 0: + return torch.nn.functional.interpolate( + input, size, scale_factor, mode, align_corners + ) + + output_shape = _output_size(2, input, size, scale_factor) + output_shape = list(input.shape[:-2]) + list(output_shape) + return _new_empty_tensor(input, output_shape) + else: + return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners) diff --git a/autonomy/il/humanoid_act/eval.py b/autonomy/il/humanoid_act/eval.py new file mode 100644 index 00000000..031c7bf9 --- /dev/null +++ b/autonomy/il/humanoid_act/eval.py @@ -0,0 +1,108 @@ +"""Not a standalone eval script — used by ``lerobot_eval.py`` for sim rollout eval.""" + +from __future__ import annotations + +from collections import deque +from pathlib import Path + +import numpy as np +import torch + +from humanoid_act.checkpoint import load_policy + + +def lerobot_cam_to_sim_name(camera_key: str) -> str: + prefix = "observation.images." + if camera_key.startswith(prefix): + return camera_key[len(prefix) :] + return camera_key + + +class LocalCustomACTPolicy: + """ + Load ``best.pt`` / ``last.pt`` from humanoid-act-train and step the sim env. + + Expects the same sim camera names as training (e.g. ego, external_D455). + State/action space matches the LeRobot dataset the checkpoint was trained on. + """ + + def __init__( + self, + robot_iface, + checkpoint_path: str, + *, + action_horizon: int | None = None, + ) -> None: + self._iface = robot_iface + self._device = torch.device(robot_iface.device) + ckpt_path = Path(checkpoint_path) + self._policy, self._config, self._stats = load_policy(ckpt_path, self._device) + self._action_horizon = action_horizon or self._config.chunk_size + self._action_queue: deque[torch.Tensor] = deque() + self._camera_order = list(self._config.camera_names) + self._validate_cameras() + print( + f"[INFO]: Loaded custom ACT checkpoint {ckpt_path} " + f"(chunk_size={self._config.chunk_size}, cameras={self._camera_order})" + ) + + def _validate_cameras(self) -> None: + sim_cams = set(self._iface.cameras.keys()) + for cam_key in self._camera_order: + sim_name = lerobot_cam_to_sim_name(cam_key) + if sim_name not in sim_cams: + raise ValueError( + f"Checkpoint expects camera '{sim_name}' (from {cam_key}) " + f"but sim only has {sorted(sim_cams)}" + ) + + def connect(self) -> None: + return + + def reset(self) -> None: + self._action_queue.clear() + + def _build_images(self, visual_obs: dict) -> np.ndarray: + images = [] + for cam_key in self._camera_order: + sim_name = lerobot_cam_to_sim_name(cam_key) + img = visual_obs[f"rgb_{sim_name}"][0].detach().cpu().numpy() + if img.dtype != np.uint8: + img = np.clip(img, 0, 255).astype(np.uint8) + if img.ndim == 3 and img.shape[-1] == 3: + img = np.transpose(img, (2, 0, 1)) + images.append(img.astype(np.float32) / 255.0) + return np.stack(images, axis=0) + + def _predict_action_chunk( + self, joint_positions: torch.Tensor, visual_obs: dict + ) -> np.ndarray: + raw_state = self._iface.get_raw_actions_from_radians(joint_positions) + qpos = self._stats.normalize_state(raw_state.cpu().numpy().astype(np.float32)) + images = self._build_images(visual_obs) + + qpos_t = torch.from_numpy(qpos).float().unsqueeze(0).to(self._device) + images_t = torch.from_numpy(images).float().unsqueeze(0).to(self._device) + + with torch.inference_mode(): + chunk = self._policy(qpos_t, images_t) + chunk_np = chunk[0].detach().cpu().numpy() + return self._stats.denormalize_action(chunk_np) + + def _refill_action_queue( + self, joint_positions: torch.Tensor, visual_obs: dict + ) -> None: + chunk = self._predict_action_chunk(joint_positions, visual_obs) + horizon = min(self._action_horizon, chunk.shape[0]) + for idx in range(horizon): + action = torch.tensor(chunk[idx], dtype=torch.float32, device=self._device) + sim_action = self._iface.get_mapped_actions_vectorized(action) + self._action_queue.append(sim_action) + + def get_action( + self, joint_positions: torch.Tensor, visual_obs: dict, log: bool = False + ) -> torch.Tensor: + del log + if not self._action_queue: + self._refill_action_queue(joint_positions, visual_obs) + return self._action_queue.popleft() diff --git a/autonomy/il/humanoid_act/normalize.py b/autonomy/il/humanoid_act/normalize.py new file mode 100644 index 00000000..f3ec4632 --- /dev/null +++ b/autonomy/il/humanoid_act/normalize.py @@ -0,0 +1,108 @@ +"""State/action normalization stats for ACT training.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np + + +STATE_KEY = "observation.state" +ACTION_KEY = "action" +MIN_STD = 1e-2 + + +@dataclass +class NormStats: + state_mean: np.ndarray + state_std: np.ndarray + action_mean: np.ndarray + action_std: np.ndarray + + def normalize_state(self, x: np.ndarray) -> np.ndarray: + return (x - self.state_mean) / self.state_std + + def normalize_action(self, x: np.ndarray) -> np.ndarray: + return (x - self.action_mean) / self.action_std + + def denormalize_state(self, x: np.ndarray) -> np.ndarray: + return x * self.state_std + self.state_mean + + def denormalize_action(self, x: np.ndarray) -> np.ndarray: + return x * self.action_std + self.action_mean + + def to_dict(self) -> dict[str, list[float]]: + return { + "state_mean": self.state_mean.tolist(), + "state_std": self.state_std.tolist(), + "action_mean": self.action_mean.tolist(), + "action_std": self.action_std.tolist(), + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> NormStats: + return cls( + state_mean=np.asarray(data["state_mean"], dtype=np.float32), + state_std=np.asarray(data["state_std"], dtype=np.float32), + action_mean=np.asarray(data["action_mean"], dtype=np.float32), + action_std=np.asarray(data["action_std"], dtype=np.float32), + ) + + +def save_stats(stats: NormStats, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(stats.to_dict(), indent=2)) + + +def load_stats(path: Path) -> NormStats: + return NormStats.from_dict(json.loads(path.read_text())) + + +def _clip_std(std: np.ndarray) -> np.ndarray: + return np.clip(std, MIN_STD, np.inf).astype(np.float32) + + +def stats_from_lerobot_meta(meta: Any) -> NormStats | None: + """Use dataset.meta.stats when LeRobot already computed them.""" + raw = getattr(meta, "stats", None) + if not raw: + return None + if STATE_KEY not in raw or ACTION_KEY not in raw: + return None + state = raw[STATE_KEY] + action = raw[ACTION_KEY] + if "mean" not in state or "std" not in state: + return None + return NormStats( + state_mean=np.asarray(state["mean"], dtype=np.float32), + state_std=_clip_std(np.asarray(state["std"], dtype=np.float32)), + action_mean=np.asarray(action["mean"], dtype=np.float32), + action_std=_clip_std(np.asarray(action["std"], dtype=np.float32)), + ) + + +def load_or_compute_stats( + dataset: Any, + *, + cache_path: Path | None = None, + recompute: bool = False, +) -> NormStats: + if cache_path and cache_path.exists() and not recompute: + return load_stats(cache_path) + + stats = stats_from_lerobot_meta(dataset.meta) + if stats is None: + raise ValueError( + "Dataset is missing meta.stats for observation.state and action. " + "Use a LeRobot dataset that includes stats (HF hub datasets do), " + "or pass a saved stats.json via cache_path." + ) + + if cache_path: + save_stats(stats, cache_path) + print(f"[stats] wrote {cache_path}") + + return stats diff --git a/autonomy/il/humanoid_act/policy.py b/autonomy/il/humanoid_act/policy.py new file mode 100644 index 00000000..0bcfc37a --- /dev/null +++ b/autonomy/il/humanoid_act/policy.py @@ -0,0 +1,70 @@ +"""ACT policy wrapper (CVAE + L1/KL loss).""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision.transforms as transforms + +from humanoid_act.config import ACTConfig +from humanoid_act.detr.build import build_act_model_and_optimizer + + +def kl_divergence(mu, logvar): + if mu.data.ndimension() == 4: + mu = mu.view(mu.size(0), mu.size(1)) + if logvar.data.ndimension() == 4: + logvar = logvar.view(logvar.size(0), logvar.size(1)) + klds = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()) + return klds.sum(1).mean(0, True) + + +class ACTPolicy(nn.Module): + def __init__(self, config: ACTConfig) -> None: + super().__init__() + self.config = config + model, optimizer = build_act_model_and_optimizer(config) + self.model = model + self._optimizer = optimizer + self.kl_weight = config.kl_weight + self._imagenet_norm = transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225], + ) + + def configure_optimizers(self) -> torch.optim.Optimizer: + return self._optimizer + + def _normalize_images(self, images: torch.Tensor) -> torch.Tensor: + b, n, c, h, w = images.shape + flat = images.reshape(b * n, c, h, w) + flat = self._imagenet_norm(flat) + return flat.reshape(b, n, c, h, w) + + def forward( + self, + qpos: torch.Tensor, + images: torch.Tensor, + actions: torch.Tensor | None = None, + is_pad: torch.Tensor | None = None, + ) -> torch.Tensor | dict[str, torch.Tensor]: + images = self._normalize_images(images) + env_state = None + + if actions is not None: + assert is_pad is not None + actions = actions[:, : self.model.num_queries] + is_pad = is_pad[:, : self.model.num_queries] + a_hat, _is_pad_hat, (mu, logvar) = self.model( + qpos, images, env_state, actions, is_pad + ) + total_kld = kl_divergence(mu, logvar) + all_l1 = F.l1_loss(actions, a_hat, reduction="none") + l1 = (all_l1 * ~is_pad.unsqueeze(-1)).mean() + kl = total_kld[0] + loss = l1 + kl * self.kl_weight + return {"loss": loss, "l1": l1, "kl": kl} + + a_hat, _, _ = self.model(qpos, images, env_state) + return a_hat diff --git a/autonomy/il/humanoid_act/smoke.py b/autonomy/il/humanoid_act/smoke.py new file mode 100644 index 00000000..188f10d4 --- /dev/null +++ b/autonomy/il/humanoid_act/smoke.py @@ -0,0 +1,79 @@ +"""Temporary smoke test for shape testing""" + +from __future__ import annotations + +import argparse +from pathlib import Path + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Smoke-test ACT dataloader on a LeRobot dataset") + parser.add_argument( + "--dataset.repo_id", + dest="repo_id", + required=True, + help="HuggingFace repo id or local repo_id (e.g. CursedRock17/so101_teleop_vials_sim_and_real)", + ) + parser.add_argument( + "--dataset.root", + dest="root", + default=None, + help="Optional local dataset root (defaults to HF cache)", + ) + parser.add_argument("--chunk_size", type=int, default=50, help="ACT action chunk length") + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument("--num_batches", type=int, default=2) + parser.add_argument( + "--stats_path", + type=Path, + default=None, + help="Optional path to write/read stats.json", + ) + args = parser.parse_args(argv) + + from humanoid_act.dataset import make_act_dataloaders + from humanoid_act.normalize import save_stats + + print(f"[smoke] repo_id={args.repo_id}") + print(f"[smoke] root={args.root or '(HF cache)'}") + print(f"[smoke] chunk_size={args.chunk_size} batch_size={args.batch_size}") + + train_loader, val_loader, stats = make_act_dataloaders( + args.repo_id, + root=args.root, + chunk_size=args.chunk_size, + batch_size=args.batch_size, + num_workers=0, + ) + + if args.stats_path: + save_stats(stats, args.stats_path) + + print( + "[smoke] norm stats:", + f"state_dim={stats.state_mean.shape[0]}", + f"action_dim={stats.action_mean.shape[0]}", + ) + print(f"[smoke] train batches ~{len(train_loader)}, val batches ~{len(val_loader)}") + + for split_name, loader in ("train", train_loader), ("val", val_loader): + print(f"\n[smoke] --- {split_name} ---") + for i, batch in enumerate(loader): + pad_frac = batch.is_pad.float().mean().item() + print( + f" batch {i}: " + f"images={tuple(batch.images.shape)} " + f"qpos={tuple(batch.qpos.shape)} " + f"actions={tuple(batch.actions.shape)} " + f"is_pad={tuple(batch.is_pad.shape)} " + f"pad_frac={pad_frac:.3f}" + ) + if i + 1 >= args.num_batches: + break + + print("\n[smoke] OK — phase 1 dataloader works.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/autonomy/il/humanoid_act/train.py b/autonomy/il/humanoid_act/train.py new file mode 100644 index 00000000..b0cb718a --- /dev/null +++ b/autonomy/il/humanoid_act/train.py @@ -0,0 +1,161 @@ +"""Train custom ACT on LeRobot datasets.""" + +from __future__ import annotations + +import argparse +import random +from dataclasses import dataclass +from itertools import cycle +from pathlib import Path + +import numpy as np +import torch + +from humanoid_act.checkpoint import save_checkpoint +from humanoid_act.config import ACTConfig +from humanoid_act.dataset import ACTBatch, make_act_dataloaders +from humanoid_act.policy import ACTPolicy + + +@dataclass +class TrainConfig: + chunk_size: int = 50 + batch_size: int = 8 + steps: int = 10_000 + val_freq: int = 500 + save_freq: int = 5000 + lr: float = 1e-5 + kl_weight: float = 10.0 + seed: int = 42 + device: str = "cuda" if torch.cuda.is_available() else "cpu" + + +def set_seed(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def forward_batch(policy: ACTPolicy, batch: ACTBatch, device: torch.device) -> dict[str, torch.Tensor]: + images = batch.images.to(device) + qpos = batch.qpos.to(device) + actions = batch.actions.to(device) + is_pad = batch.is_pad.to(device) + return policy(qpos, images, actions, is_pad) + + +@torch.no_grad() +def validate(policy: ACTPolicy, val_loader, device: torch.device) -> dict[str, float]: + policy.eval() + totals: dict[str, float] = {} + count = 0 + for batch in val_loader: + loss_dict = forward_batch(policy, batch, device) + count += 1 + for key, value in loss_dict.items(): + totals[key] = totals.get(key, 0.0) + float(value.item()) + return {key: value / max(count, 1) for key, value in totals.items()} + + +def main(argv: list[str] | None = None) -> int: + train_cfg = TrainConfig() + + parser = argparse.ArgumentParser(description="Train humanoid ACT on a LeRobot dataset") + parser.add_argument("repo_id", help="HuggingFace repo id or local LeRobot repo_id") + parser.add_argument("output_dir", type=Path, help="Directory for checkpoints and config.json") + parser.add_argument( + "--root", + default=None, + help="Optional local dataset root (defaults to HF cache)", + ) + args = parser.parse_args(argv) + + set_seed(train_cfg.seed) + device = torch.device(train_cfg.device) + args.output_dir.mkdir(parents=True, exist_ok=True) + + train_loader, val_loader, stats = make_act_dataloaders( + args.repo_id, + root=args.root, + chunk_size=train_cfg.chunk_size, + batch_size=train_cfg.batch_size, + seed=train_cfg.seed, + ) + + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + meta = LeRobotDataset(args.repo_id, root=args.root).meta + config = ACTConfig.from_dataset_meta( + args.repo_id, + meta, + chunk_size=train_cfg.chunk_size, + kl_weight=train_cfg.kl_weight, + lr=train_cfg.lr, + seed=train_cfg.seed, + ) + config.save(args.output_dir / "config.json") + + policy = ACTPolicy(config).to(device) + optimizer = policy.configure_optimizers() + train_iter = cycle(train_loader) + + best_val = float("inf") + print( + f"[train] repo_id={args.repo_id} device={device} " + f"state_dim={config.state_dim} chunk_size={config.chunk_size} " + f"cameras={config.camera_names}" + ) + + for step in range(1, train_cfg.steps + 1): + policy.train() + batch = next(train_iter) + loss_dict = forward_batch(policy, batch, device) + loss = loss_dict["loss"] + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + if step % 100 == 0 or step == 1: + print( + f"[train] step {step}/{train_cfg.steps} " + f"loss={loss.item():.4f} l1={loss_dict['l1'].item():.4f} " + f"kl={loss_dict['kl'].item():.4f}" + ) + + if step % train_cfg.val_freq == 0 or step == train_cfg.steps: + val_metrics = validate(policy, val_loader, device) + val_loss = val_metrics["loss"] + print( + f"[val] step {step} loss={val_loss:.4f} " + f"l1={val_metrics['l1']:.4f} kl={val_metrics['kl']:.4f}" + ) + if val_loss < best_val: + best_val = val_loss + save_checkpoint( + args.output_dir / "best.pt", + policy=policy, + config=config, + stats=stats, + step=step, + val_loss=val_loss, + ) + print(f"[train] saved best checkpoint (val_loss={val_loss:.4f})") + + if step % train_cfg.save_freq == 0 or step == train_cfg.steps: + save_checkpoint( + args.output_dir / "last.pt", + policy=policy, + config=config, + stats=stats, + step=step, + ) + print(f"[train] saved last checkpoint at step {step}") + + print(f"[train] done. best_val_loss={best_val:.4f}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/autonomy/il/pyproject.toml b/autonomy/il/pyproject.toml index d045b6ec..6562e661 100644 --- a/autonomy/il/pyproject.toml +++ b/autonomy/il/pyproject.toml @@ -23,7 +23,9 @@ sim = ["torch>=2.0", "tqdm>=4.0", "lerobot>=0.3.0"] [project.scripts] humanoid-record = "humanoid_il.record:main" +humanoid-act-smoke = "humanoid_act.smoke:main" +humanoid-act-train = "humanoid_act.train:main" [tool.setuptools.packages.find] where = ["."] -include = ["humanoid_il*"] +include = ["humanoid_il*", "humanoid_act*"] diff --git a/autonomy/simulation/so101_vial_task/scripts/lerobot_eval.py b/autonomy/simulation/so101_vial_task/scripts/lerobot_eval.py index 21d742f0..f3eba50b 100644 --- a/autonomy/simulation/so101_vial_task/scripts/lerobot_eval.py +++ b/autonomy/simulation/so101_vial_task/scripts/lerobot_eval.py @@ -41,15 +41,18 @@ parser.add_argument( "--policy_type", type=str, - choices=("lerobot", "groot"), + choices=("lerobot", "groot", "custom_act"), default="lerobot", - help="lerobot: local checkpoint; groot: remote GR00T server", + help="lerobot: LeRobot checkpoint dir; custom_act: humanoid-act best.pt; groot: GR00T server", ) parser.add_argument( "--policy_path", type=str, default=None, - help="Local LeRobot policy checkpoint dir (required for --policy_type lerobot)", + help=( + "Policy checkpoint path: LeRobot pretrained_model dir (--policy_type lerobot), " + "or humanoid-act .pt file (--policy_type custom_act)" + ), ) parser.add_argument( "--rename_map", @@ -97,6 +100,7 @@ import humanoid_so101_vial_task.tasks # noqa: F401 from humanoid_so101_vial_task.utils.keyboard import KeyboardControl +from humanoid_act.eval import LocalCustomACTPolicy from humanoid_so101_vial_task.utils.lerobot_interface import ( LeRobotSO101Interface, GR00TRemotePolicy, @@ -168,6 +172,17 @@ def main(): lang_instruction=args_cli.lang_instruction, ) policy.connect() + elif args_cli.policy_type == "custom_act": + if not args_cli.policy_path: + raise ValueError("--policy_path is required when --policy_type custom_act") + if args_cli.rename_map: + print("[WARNING] --rename_map is ignored for custom_act (use training camera names)") + policy = LocalCustomACTPolicy( + robot_iface=robot_iface, + checkpoint_path=args_cli.policy_path, + action_horizon=args_cli.action_horizon, + ) + policy.connect() else: if not args_cli.policy_path: raise ValueError("--policy_path is required when --policy_type lerobot")