| """Apply the trained command policy and write a physics-show timeline.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnxruntime as ort |
|
|
|
|
| def run(args: argparse.Namespace) -> dict: |
| source = np.load(args.retarget.resolve(), allow_pickle=True) |
| session = ort.InferenceSession(str(args.policy.resolve())) |
| commands = session.run( |
| [session.get_outputs()[0].name], |
| {session.get_inputs()[0].name: source["human_features_14d"].astype(np.float32)}, |
| )[0].astype(np.float32) |
| rule_commands = source["command_13d"].astype(np.float32) |
| |
| |
| |
| commands[:, 0:3] = rule_commands[:, 0:3] |
| |
| |
| commands[:, 0] = np.clip(commands[:, 0], 0.0, 0.30) |
| commands[:, 1] = 0.0 |
| commands[:, 2] = np.clip(commands[:, 2], -1.0, 1.0) |
| commands[:, 3:7] = np.clip(commands[:, 3:7], -0.20, 0.20) |
| commands[:, 7] = 0.0 |
| commands[:, 8] = 0.0 |
| commands[:, 9] = np.clip(commands[:, 9], -0.015, 0.005) |
| commands[:, 10:12] = np.clip(commands[:, 10:12], -0.12, 0.12) |
| commands[:, 12] = 0.0 |
| output = args.output.resolve() |
| output.parent.mkdir(parents=True, exist_ok=True) |
| payload = {name: source[name] for name in source.files} |
| payload["rule_based_command_13d"] = payload.pop("command_13d") |
| payload["command_13d"] = commands |
| payload["command_policy_onnx"] = np.asarray(str(args.policy)) |
| np.savez_compressed(output, **payload) |
| error = commands - source["command_13d"] |
| summary = { |
| "schema_version": "1.0", |
| "policy": str(args.policy), |
| "source_retarget": str(args.retarget), |
| "output": str(output), |
| "rmse_command": float(np.sqrt(np.mean(error ** 2))), |
| "max_abs_error_by_command": np.max(np.abs(error), axis=0).round(7).tolist(), |
| "hard_limits_applied": True, |
| "safety_channel_lock": { |
| "rule_based_validated": "twist[0:3] and kick-policy schedule", |
| "learned": "head_pose[3:7] and body_pose[7:13]", |
| }, |
| } |
| output.with_suffix(".application.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("--policy", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| run(parser.parse_args()) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
