Files
RC_WheelLeg/05_software/real/sim2real/tools/alignment_check.py
T

172 lines
6.8 KiB
Python

"""Offline deployment alignment check for the current 53-D rough policy."""
import argparse
from pathlib import Path
import sys
import numpy as np
import torch
import yaml
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from interface.motor_mapping import MotorMapping # noqa: E402
from policy.policy_runner import PolicyRunner # noqa: E402
def _load_manifest(manifest_path: Path) -> dict:
with open(manifest_path, "r", encoding="utf-8") as f:
return yaml.safe_load(f)
def check(policy_path: Path, manifest_path: Path | None = None) -> int:
issues: list[tuple[str, str]] = []
manifest = _load_manifest(manifest_path) if manifest_path is not None else None
expected = (
("fl", "hip_abduction"), ("fl", "hip_pitch"), ("fl", "knee"),
("fr", "hip_abduction"), ("fr", "hip_pitch"), ("fr", "knee"),
("rl", "hip_abduction"), ("rl", "hip_pitch"), ("rl", "knee"),
("rr", "hip_abduction"), ("rr", "hip_pitch"), ("rr", "knee"),
("fl", "wheel"), ("fr", "wheel"), ("rl", "wheel"), ("rr", "wheel"),
)
if MotorMapping.SIM_JOINT_ORDER != expected:
issues.append(("joint_order", "MotorMapping.SIM_JOINT_ORDER mismatch"))
else:
print("[Check] joint order: PASS")
manifest_enable_zero_cmd = True
if manifest is not None:
manifest_enable_zero_cmd = bool(
manifest.get("model", {}).get("enable_zero_cmd_suppression", True)
)
runner = PolicyRunner(
policy_path,
device=torch.device("cpu"),
enable_zero_cmd_suppression=manifest_enable_zero_cmd,
)
obs_mean = runner.policy.obs_mean.detach().cpu().numpy()
obs_std = runner.policy.obs_std.detach().cpu().numpy()
if np.allclose(obs_mean, 0.0) and np.allclose(obs_std, 1.0):
print("[Check] obs normalizer: PASS (identity)")
else:
print(
f"[Check] obs normalizer: PASS "
f"(mean range=[{obs_mean.min():.3f},{obs_mean.max():.3f}], "
f"std range=[{obs_std.min():.3f},{obs_std.max():.3f}])"
)
if (obs_std < 1e-6).any():
issues.append(
(
"normalizer_zero_std",
f"obs_std has near-zero entries: {np.where(obs_std < 1e-6)[0].tolist()}",
)
)
raw_zero = np.zeros(runner.BASE_OBS_DIM, dtype=np.float32)
runner.reset(prime_obs=raw_zero)
_, raw = runner.step(raw_zero)
if np.max(np.abs(raw)) > 5.0:
issues.append(
(
"output_range",
f"raw action too large under zero obs: {np.max(np.abs(raw)):.3f}",
)
)
else:
print(f"[Check] zero-obs output range: PASS (max|raw|={np.max(np.abs(raw)):.3f})")
expected_default = np.array([0.0, 0.9, -1.8] * 4 + [0.0] * 4, dtype=np.float32)
if not np.allclose(runner.default_dof_pos, expected_default):
issues.append(("default_pose_mismatch", f"default_dof_pos mismatch: {runner.default_dof_pos}"))
else:
print("[Check] default_dof_pos: PASS")
if runner.BASE_OBS_DIM != 53:
issues.append(("obs_dim", f"base obs dim {runner.BASE_OBS_DIM} != 53"))
else:
print("[Check] actor obs dim: PASS (53)")
if manifest is not None:
declared_model = manifest.get("model", {})
declared_action = manifest.get("action", {})
declared_safety = manifest.get("safety", {})
declared_control = manifest.get("control", {})
if int(declared_model.get("obs_dim", -1)) != runner.policy.expected_obs_dim:
issues.append(
(
"manifest_obs_dim",
f"manifest obs_dim {declared_model.get('obs_dim')} != policy {runner.policy.expected_obs_dim}",
)
)
else:
print("[Check] manifest obs_dim: PASS")
if int(declared_model.get("action_dim", -1)) != runner.policy.expected_action_dim:
issues.append(
(
"manifest_action_dim",
f"manifest action_dim {declared_model.get('action_dim')} != policy {runner.policy.expected_action_dim}",
)
)
else:
print("[Check] manifest action_dim: PASS")
declared_default = np.asarray(declared_action.get("default_dof_pos", []), dtype=np.float32)
if declared_default.shape != runner.default_dof_pos.shape or not np.allclose(
declared_default, runner.default_dof_pos
):
issues.append(("manifest_default_pose", "manifest default_dof_pos mismatch"))
else:
print("[Check] manifest default_dof_pos: PASS")
declared_scale = np.asarray(declared_action.get("scale", []), dtype=np.float32)
if declared_scale.shape != runner.action_scale.shape or not np.allclose(
declared_scale, runner.action_scale
):
issues.append(("manifest_action_scale", "manifest action scale mismatch"))
else:
print("[Check] manifest action scale: PASS")
if float(declared_safety.get("zero_cmd_lin_thresh", -1.0)) != runner.zero_cmd_lin_thresh:
issues.append(("manifest_zero_cmd_lin_thresh", "manifest zero_cmd_lin_thresh mismatch"))
if float(declared_safety.get("zero_cmd_yaw_thresh", -1.0)) != runner.zero_cmd_yaw_thresh:
issues.append(("manifest_zero_cmd_yaw_thresh", "manifest zero_cmd_yaw_thresh mismatch"))
if float(declared_safety.get("zero_yaw_rate_thresh", -1.0)) != runner.zero_yaw_rate_thresh:
issues.append(("manifest_zero_yaw_rate_thresh", "manifest zero_yaw_rate_thresh mismatch"))
if bool(declared_model.get("enable_zero_cmd_suppression", True)) != runner.enable_zero_cmd_suppression:
issues.append(("manifest_zero_cmd_switch", "manifest zero-command suppression switch mismatch"))
else:
print("[Check] manifest zero-command suppression: PASS")
if int(declared_control.get("control_freq_hz", -1)) != 50:
issues.append(("manifest_control_freq", "manifest control_freq_hz must be 50"))
else:
print("[Check] manifest control freq: PASS")
if issues:
print("\n" + "=" * 60)
print(f"Alignment check failed: {len(issues)} issue(s)")
for tag, msg in issues:
print(f" [{tag}] {msg}")
return 1
print("\n" + "=" * 60)
print("All offline alignment checks passed.")
return 0
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--policy", type=str, required=True, help="Path to policy .pt")
parser.add_argument("--manifest", type=str, default=None, help="Optional deployment manifest yaml")
args = parser.parse_args()
manifest = Path(args.manifest) if args.manifest else None
sys.exit(check(Path(args.policy), manifest))
if __name__ == "__main__":
main()