File size: 2,899 Bytes
32502a8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """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)
# Safety-critical gait velocity and turn timing stay on the already passed
# human-derived reference. The learned policy supplies the expressive
# head/body channels; the official Pollen policy supplies joint balance.
commands[:, 0:3] = rule_commands[:, 0:3]
# Hard deployment envelope. Preserve intentionally unsupported lateral
# command and body-translation axes at exact zero.
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()
|