fix: clamp std before distribution, lower init_noise to 0.5, NaN guard
This commit is contained in:
8x54zj-m
289
scripts/export_go1_no_linevel_onnx.py
Normal file
289
scripts/export_go1_no_linevel_onnx.py
Normal file
@@ -0,0 +1,289 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Export JAX/Flax-trained SKRL Go1 (no-linevel, 57-dim) policy to ONNX.
|
||||
|
||||
Converts Flax weights → PyTorch → ONNX, baking in the RunningStandardScaler
|
||||
normalization so the ONNX model accepts raw (scaled) observations directly.
|
||||
|
||||
Usage:
|
||||
uv run scripts/export_go1_no_linevel_onnx.py
|
||||
uv run scripts/export_go1_no_linevel_onnx.py --output ./my_exports
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
|
||||
import msgpack
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _decode_flax_array(ext) -> np.ndarray | None:
|
||||
if not hasattr(ext, "code"):
|
||||
return None
|
||||
parts = msgpack.unpackb(ext.data, raw=False)
|
||||
if not isinstance(parts, list) or len(parts) < 3:
|
||||
return None
|
||||
|
||||
def _flatten(s):
|
||||
if isinstance(s, list):
|
||||
out = []
|
||||
for item in s:
|
||||
out.extend(_flatten(item))
|
||||
return out
|
||||
return [s]
|
||||
|
||||
shape = tuple(_flatten(parts[0]))
|
||||
dtype_str = parts[1]
|
||||
if isinstance(dtype_str, bytes):
|
||||
dtype_str = dtype_str.decode("utf-8")
|
||||
raw = parts[2]
|
||||
return np.frombuffer(raw, dtype=np.dtype(dtype_str)).reshape(shape)
|
||||
|
||||
|
||||
def load_jax_checkpoint(ckpt_path: str) -> dict:
|
||||
with open(ckpt_path, "rb") as f:
|
||||
data = pickle.load(f)
|
||||
|
||||
policy_raw = msgpack.unpackb(data["policy"])
|
||||
flax_params = {}
|
||||
for name, val in policy_raw["params"].items():
|
||||
if isinstance(val, dict):
|
||||
flax_params[name] = {k: _decode_flax_array(v) for k, v in val.items()}
|
||||
else:
|
||||
flax_params[name] = _decode_flax_array(val)
|
||||
|
||||
prep = msgpack.unpackb(data["state_preprocessor"])
|
||||
running_mean = _decode_flax_array(prep["running_mean"])
|
||||
running_var = _decode_flax_array(prep["running_variance"])
|
||||
count_arr = _decode_flax_array(prep["current_count"])
|
||||
count = int(count_arr.flat[0]) if count_arr is not None else 0
|
||||
|
||||
return {
|
||||
"flax_params": flax_params,
|
||||
"running_mean": running_mean,
|
||||
"running_var": running_var,
|
||||
"count": count,
|
||||
}
|
||||
|
||||
|
||||
# -- PyTorch model -----------------------------------------------------------
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class PolicyTorch(nn.Module):
|
||||
def __init__(self, obs_dim: int, action_dim: int, hidden_dims: list[int]):
|
||||
super().__init__()
|
||||
self.obs_dim = obs_dim
|
||||
self.action_dim = action_dim
|
||||
self.hidden_dims = hidden_dims
|
||||
|
||||
layers = []
|
||||
in_dim = obs_dim
|
||||
for h in hidden_dims:
|
||||
layers.extend([nn.Linear(in_dim, h), nn.ELU()])
|
||||
in_dim = h
|
||||
self.net = nn.Sequential(*layers)
|
||||
self.mean_layer = nn.Linear(in_dim, action_dim)
|
||||
|
||||
def forward(self, x):
|
||||
return self.mean_layer(self.net(x))
|
||||
|
||||
|
||||
class ONNXExporter(nn.Module):
|
||||
def __init__(self, policy: PolicyTorch, mean: np.ndarray, std: np.ndarray):
|
||||
super().__init__()
|
||||
self.policy = policy
|
||||
self.register_buffer("mean", torch.from_numpy(mean).float())
|
||||
self.register_buffer("std", torch.from_numpy(std).float())
|
||||
self.clip_threshold = 5.0
|
||||
|
||||
def forward(self, x):
|
||||
x = (x - self.mean) / (self.std + 1e-8)
|
||||
x = torch.clamp(x, min=-self.clip_threshold, max=self.clip_threshold)
|
||||
return self.policy(x)
|
||||
|
||||
|
||||
def flax_to_torch_weights(flax_params: dict, obs_dim: int, hidden_dims: list[int], action_dim: int) -> dict:
|
||||
state_dict = {}
|
||||
layer_names = sorted([k for k in flax_params if k.startswith("Dense_")])
|
||||
|
||||
hidden_dense = layer_names[:-1]
|
||||
layer_idx = 0
|
||||
for name in hidden_dense:
|
||||
layer_params = flax_params[name]
|
||||
kernel = layer_params["kernel"]
|
||||
bias = layer_params["bias"]
|
||||
state_dict[f"net.{layer_idx}.weight"] = torch.from_numpy(kernel.T.copy()).float()
|
||||
state_dict[f"net.{layer_idx}.bias"] = torch.from_numpy(bias.copy()).float()
|
||||
layer_idx += 2
|
||||
|
||||
last_name = layer_names[-1]
|
||||
last_params = flax_params[last_name]
|
||||
state_dict["mean_layer.weight"] = torch.from_numpy(last_params["kernel"].T.copy()).float()
|
||||
state_dict["mean_layer.bias"] = torch.from_numpy(last_params["bias"].copy()).float()
|
||||
|
||||
return state_dict
|
||||
|
||||
|
||||
# -- Config ------------------------------------------------------------------
|
||||
|
||||
GO1_JOINT_NAMES = [
|
||||
"FR_hip", "FR_thigh", "FR_calf",
|
||||
"FL_hip", "FL_thigh", "FL_calf",
|
||||
"RR_hip", "RR_thigh", "RR_calf",
|
||||
"RL_hip", "RL_thigh", "RL_calf",
|
||||
]
|
||||
|
||||
GO1_DEFAULT_ANGLES = np.array([
|
||||
-0.0, 0.9, -1.8,
|
||||
0.0, 0.9, -1.8,
|
||||
-0.0, 0.9, -1.8,
|
||||
0.0, 0.9, -1.8,
|
||||
], dtype=np.float32)
|
||||
|
||||
# 57-dim observation layout (NO linear velocity):
|
||||
# [0:3] gyro (scaled *0.25)
|
||||
# [3:6] gravity vector (body frame)
|
||||
# [6:18] joint angle deviation (scaled *1.0)
|
||||
# [18:30] joint velocity (scaled *0.05)
|
||||
# [30:42] last action (raw)
|
||||
# [42:45] command [vx*2.0, vy*2.0, wz*0.25]
|
||||
# [45:57] foot contact forces (body frame, raw)
|
||||
|
||||
OBS_SCALES = {
|
||||
"ang_vel": 0.25,
|
||||
"dof_pos": 1.0,
|
||||
"dof_vel": 0.05,
|
||||
"contact_force": 1.0, # raw, no scaling
|
||||
}
|
||||
|
||||
ACTION_SCALE = 0.05
|
||||
KP, KD = 80.0, 1.0
|
||||
CLIP_ACTIONS = 23.7
|
||||
CLIP_OBS = 100.0
|
||||
|
||||
|
||||
def auto_discover_checkpoint(env_name: str) -> str:
|
||||
"""Find the latest best_agent checkpoint for the given env."""
|
||||
base_dir = Path(f"runs/{env_name}/skrl")
|
||||
if not base_dir.exists():
|
||||
raise FileNotFoundError(f"No training results found: {base_dir}")
|
||||
|
||||
runs = sorted([d for d in base_dir.iterdir() if d.is_dir()], key=lambda d: d.stat().st_mtime, reverse=True)
|
||||
for run_dir in runs:
|
||||
ckpt = run_dir / "checkpoints" / "best_agent.pickle"
|
||||
if ckpt.exists():
|
||||
return str(ckpt)
|
||||
|
||||
raise FileNotFoundError(f"No best_agent.pickle found in {base_dir}")
|
||||
|
||||
|
||||
def export(checkpoint_path: str, output_dir: str):
|
||||
ckpt = load_jax_checkpoint(checkpoint_path)
|
||||
flax_params = ckpt["flax_params"]
|
||||
running_mean = ckpt["running_mean"]
|
||||
running_var = ckpt["running_var"]
|
||||
running_std = np.sqrt(running_var)
|
||||
|
||||
dense_keys = sorted([k for k in flax_params if k.startswith("Dense_")])
|
||||
hidden_dims = [flax_params[k]["bias"].shape[0] for k in dense_keys[:-1]]
|
||||
obs_dim = flax_params[dense_keys[0]]["kernel"].shape[0]
|
||||
action_dim = flax_params[dense_keys[-1]]["bias"].shape[0]
|
||||
|
||||
print(f"Architecture: obs={obs_dim}, hidden={hidden_dims}, action={action_dim}")
|
||||
print(f"Normalizer mean range: [{running_mean.min():.4f}, {running_mean.max():.4f}]")
|
||||
print(f"Normalizer std range: [{running_std.min():.6f}, {running_std.max():.6f}]")
|
||||
|
||||
policy = PolicyTorch(obs_dim, action_dim, hidden_dims)
|
||||
torch_weights = flax_to_torch_weights(flax_params, obs_dim, hidden_dims, action_dim)
|
||||
policy.load_state_dict(torch_weights, strict=True)
|
||||
policy.eval()
|
||||
|
||||
rng = np.random.RandomState(42)
|
||||
test_obs = rng.randn(1, obs_dim).astype(np.float32)
|
||||
with torch.no_grad():
|
||||
torch_out = policy(torch.from_numpy(test_obs)).numpy()
|
||||
print(f"Test forward pass: input shape={test_obs.shape}, output shape={torch_out.shape}")
|
||||
print(f" output sample: {np.array2string(torch_out[0, :4], precision=4, suppress_small=True)} ...")
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
onnx_path = os.path.join(output_dir, "policy.onnx")
|
||||
|
||||
exporter = ONNXExporter(policy, running_mean, running_std)
|
||||
exporter.eval()
|
||||
|
||||
dummy = torch.zeros(1, obs_dim, dtype=torch.float32)
|
||||
torch.onnx.export(
|
||||
exporter, dummy, onnx_path,
|
||||
export_params=True, opset_version=11,
|
||||
input_names=["observations"], output_names=["actions"],
|
||||
dynamic_axes={},
|
||||
)
|
||||
print(f"ONNX exported to: {onnx_path}")
|
||||
|
||||
# Normalizer stats
|
||||
npz_path = os.path.join(output_dir, "normalizer.npz")
|
||||
np.savez(npz_path, mean=running_mean, std=running_std)
|
||||
print(f"Normalizer saved to: {npz_path}")
|
||||
|
||||
# Metadata
|
||||
meta_path = os.path.join(output_dir, "metadata.txt")
|
||||
with open(meta_path, "w") as f:
|
||||
f.write("# Go1 No-Linevel Terrain Walk - ONNX Policy Metadata\n")
|
||||
f.write(f"env: go1-stairs-terrain-walk-no-linevel\n")
|
||||
f.write(f"obs_dim: {obs_dim}\n")
|
||||
f.write(f"action_dim: {action_dim}\n")
|
||||
f.write(f"hidden_dims: {hidden_dims}\n")
|
||||
f.write("\n# Observation layout (57 dims, NO linear velocity):\n")
|
||||
f.write(" [0:3] gyro * ang_vel_scale\n")
|
||||
f.write(" [3:6] gravity (body frame)\n")
|
||||
f.write(" [6:18] joint_angle_deviation * dof_pos_scale\n")
|
||||
f.write(" [18:30] joint_vel * dof_vel_scale\n")
|
||||
f.write(" [30:42] last_actions (raw)\n")
|
||||
f.write(" [42:45] commands [vx*2.0, vy*2.0, wz*0.25]\n")
|
||||
f.write(" [45:57] foot_contact_forces (body frame, raw)\n")
|
||||
f.write(f"\n# Joint order: {GO1_JOINT_NAMES}\n")
|
||||
f.write(f"default_angles: {GO1_DEFAULT_ANGLES.tolist()}\n")
|
||||
f.write(f"action_scale: {ACTION_SCALE}\n")
|
||||
f.write(f"kp: {KP}\n")
|
||||
f.write(f"kd: {KD}\n")
|
||||
f.write(f"clip_actions: {CLIP_ACTIONS}\n")
|
||||
f.write(f"clip_observations: {CLIP_OBS}\n")
|
||||
f.write("\n# Observation scales:\n")
|
||||
for k, v in OBS_SCALES.items():
|
||||
f.write(f" {k}: {v}\n")
|
||||
f.write(" command_scale: [2.0, 2.0, 0.25]\n")
|
||||
print(f"Metadata saved to: {meta_path}")
|
||||
|
||||
return onnx_path
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Export JAX-trained Go1 no-linevel policy to ONNX")
|
||||
parser.add_argument("--checkpoint", type=str, default=None,
|
||||
help="Path to SKRL JAX checkpoint (auto-discovered if not set)")
|
||||
parser.add_argument("--output", type=str, default="exports_go1_no_linevel",
|
||||
help="Output directory")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.checkpoint:
|
||||
ckpt_path = args.checkpoint
|
||||
else:
|
||||
ckpt_path = auto_discover_checkpoint("go1-stairs-terrain-walk-no-linevel")
|
||||
|
||||
if not os.path.exists(ckpt_path):
|
||||
print(f"Error: checkpoint not found: {ckpt_path}")
|
||||
return 1
|
||||
|
||||
print(f"Loading checkpoint: {ckpt_path}")
|
||||
onnx_path = export(ckpt_path, args.output)
|
||||
print(f"\nDone! ONNX model ready for sim2sim:")
|
||||
print(f" {onnx_path}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit(main())
|
||||
Reference in New Issue
Block a user