"""Train the stable human-motion-to-Microduck-command policy used by the show.""" from __future__ import annotations import argparse import csv import json from pathlib import Path import random import time import numpy as np import onnx import torch from torch import nn from torch.utils.data import DataLoader, TensorDataset class CommandPolicy(nn.Module): def __init__( self, input_dim: int, feature_mean: torch.Tensor, feature_std: torch.Tensor, command_mean: torch.Tensor, command_std: torch.Tensor, ) -> None: super().__init__() self.register_buffer("feature_mean", feature_mean) self.register_buffer("feature_std", feature_std) self.register_buffer("command_mean", command_mean) self.register_buffer("command_std", command_std) self.net = nn.Sequential( nn.Linear(input_dim, 128), nn.ELU(), nn.Linear(128, 128), nn.ELU(), nn.Linear(128, 64), nn.ELU(), nn.Linear(64, 13), ) def forward(self, features: torch.Tensor) -> torch.Tensor: normalized = (features - self.feature_mean) / self.feature_std return self.net(normalized) * self.command_std + self.command_mean def normalized_command(self, features: torch.Tensor) -> torch.Tensor: normalized = (features - self.feature_mean) / self.feature_std return self.net(normalized) def run(args: argparse.Namespace) -> dict: random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) started = time.time() source = np.load(args.retarget.resolve(), allow_pickle=True) features = source["human_features_14d"].astype(np.float32) targets = source["command_13d"].astype(np.float32) names = [str(value) for value in source["human_feature_names"]] val_mask = np.arange(len(features)) % 5 == 0 train_mask = ~val_mask x_train, y_train = features[train_mask], targets[train_mask] x_val, y_val = features[val_mask], targets[val_mask] feature_mean_np = x_train.mean(axis=0) feature_std_np = np.maximum(x_train.std(axis=0), 1e-4) command_mean_np = y_train.mean(axis=0) command_std_np = np.maximum(y_train.std(axis=0), 1e-4) feature_mean = torch.tensor(feature_mean_np) feature_std = torch.tensor(feature_std_np) command_mean = torch.tensor(command_mean_np) command_std = torch.tensor(command_std_np) input_dim = features.shape[1] model = CommandPolicy(input_dim, feature_mean, feature_std, command_mean, command_std) train_y_normalized = (y_train - command_mean_np) / command_std_np loader = DataLoader( TensorDataset(torch.tensor(x_train), torch.tensor(train_y_normalized)), batch_size=args.batch_size, shuffle=True, generator=torch.Generator().manual_seed(args.seed), ) optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=1e-6) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) noise_std = torch.tensor( [0.001, 0.001, 0.005, 0.005, 0.005, 0.015, 0.015, 0.01, 0.01, 0.005, 0.005, 0.0, 0.0, 0.0] ) loss_fn = nn.SmoothL1Loss(beta=0.08) x_val_t = torch.tensor(x_val) y_val_t = torch.tensor(y_val) rows: list[dict] = [] best_rmse = float("inf") best_state = None for epoch in range(1, args.epochs + 1): model.train() total = 0.0 count = 0 for batch_x, batch_y in loader: augmented = batch_x + torch.randn_like(batch_x) * noise_std predicted = model.normalized_command(augmented) loss = loss_fn(predicted, batch_y) optimizer.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total += float(loss.detach()) * len(batch_x) count += len(batch_x) scheduler.step() model.eval() with torch.no_grad(): predicted = model(x_val_t) rmse = torch.sqrt(torch.mean((predicted - y_val_t) ** 2)) mae = torch.mean(torch.abs(predicted - y_val_t)) row = { "epoch": epoch, "train_smooth_l1": total / max(1, count), "val_rmse_command": float(rmse), "val_mae_command": float(mae), "learning_rate": optimizer.param_groups[0]["lr"], } rows.append(row) if row["val_rmse_command"] < best_rmse: best_rmse = row["val_rmse_command"] best_state = { key: value.detach().cpu().clone() for key, value in model.state_dict().items() } if epoch == 1 or epoch % 50 == 0 or epoch == args.epochs: print(json.dumps(row)) assert best_state is not None model.load_state_dict(best_state) model.eval() output = args.output.resolve() output.mkdir(parents=True, exist_ok=True) torch.save({ "state_dict": best_state, "seed": args.seed, "feature_names": names, "architecture": [input_dim, 128, 128, 64, 13], }, output / "human_motion_command_policy.pt") with (output / "training_metrics.csv").open("w", newline="", encoding="utf-8") as stream: writer = csv.DictWriter(stream, fieldnames=rows[0].keys()) writer.writeheader() writer.writerows(rows) np.savez_compressed( output / "normalization.npz", feature_mean=feature_mean_np, feature_std=feature_std_np, command_mean=command_mean_np, command_std=command_std_np, feature_names=np.asarray(names), ) onnx_path = output / "human_motion_command_policy.onnx" torch.onnx.export( model, torch.zeros((1, input_dim), dtype=torch.float32), str(onnx_path), input_names=["human_motion_features"], output_names=["microduck_command_13d"], dynamic_axes={ "human_motion_features": {0: "batch"}, "microduck_command_13d": {0: "batch"}, }, opset_version=17, dynamo=False, ) graph = onnx.load(str(onnx_path)) for key, value in { "input_features": ",".join(names), "output_contract": "twist(3),head_pose(4),body_pose(6)", "control_hz": "50", "target_bpm": "111", "low_level_controller": "official Pollen Microduck policy family", }.items(): item = graph.metadata_props.add() item.key = key item.value = value onnx.save(graph, str(onnx_path)) with torch.no_grad(): full_pred = model(torch.tensor(features)).numpy() error = full_pred - targets summary = { "schema_version": "1.0", "training_method": "supervised human-motion feature to deployable command policy", "seed": args.seed, "epochs": args.epochs, "frames": len(features), "best_validation_rmse_command": best_rmse, "full_sequence_rmse_command": float(np.sqrt(np.mean(error ** 2))), "max_abs_command_error": np.max(np.abs(error), axis=0).round(7).tolist(), "elapsed_seconds": time.time() - started, "input_features": names, "output_contract": "[twist(3), head_pose(4), body_pose(6)]", "safety_architecture": "learned choreography command policy + official Pollen balancing policies", "files": { "checkpoint": "human_motion_command_policy.pt", "onnx": "human_motion_command_policy.onnx", "metrics": "training_metrics.csv", "normalization": "normalization.npz", }, } (output / "training_summary.json").write_text( json.dumps(summary, indent=2), encoding="utf-8", ) print(json.dumps(summary, indent=2)) return summary def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--retarget", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--epochs", type=int, default=500) parser.add_argument("--batch-size", type=int, default=256) parser.add_argument("--learning-rate", type=float, default=1e-3) parser.add_argument("--seed", type=int, default=20260830) run(parser.parse_args()) if __name__ == "__main__": main()