| """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() |
|
|
