218 lines
9.5 KiB
Python
218 lines
9.5 KiB
Python
#!/usr/bin/env python3
|
|
"""DreamWaQ ONNX export: CENet encoder + Actor → ONNX for MuJoCo deployment.
|
|
|
|
Model: 2 inputs, 1 output
|
|
- observations: (1, 45)
|
|
- obs_history: (1, 5, 45)
|
|
→ actions: (1, 12)
|
|
|
|
Usage:
|
|
uv run scripts/export_dreamwaq_onnx.py
|
|
uv run scripts/export_dreamwaq_onnx.py --checkpoint PATH --vae PATH --output PATH
|
|
"""
|
|
import argparse, os, pickle, sys
|
|
import msgpack
|
|
import numpy as np
|
|
import jax, jax.numpy as jnp
|
|
import flax.linen as nn
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# Export helpers
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
def _decode_flax_array(ext) -> np.ndarray | None:
|
|
"""Decode flax-serialized msgpack ExtType to numpy array."""
|
|
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
|
|
shape = []
|
|
def _flatten(s):
|
|
if isinstance(s, list):
|
|
for x in s: _flatten(x)
|
|
elif isinstance(s, int):
|
|
shape.append(s)
|
|
_flatten(parts[0])
|
|
dtype_str = parts[1]
|
|
raw_bytes = parts[2]
|
|
return np.frombuffer(raw_bytes, dtype=np.dtype(dtype_str)).reshape(shape)
|
|
|
|
|
|
def load_skrl_policy(path):
|
|
"""Extract actor weights from SKRL checkpoint. Slice first layer: 254→64."""
|
|
with open(path, 'rb') as f:
|
|
ckpt = pickle.load(f)
|
|
raw = msgpack.unpackb(ckpt['policy'])['params']
|
|
params = {}
|
|
for name, val in raw.items():
|
|
if isinstance(val, dict):
|
|
params[name] = {k: _decode_flax_array(v) for k, v in val.items()}
|
|
else:
|
|
params[name] = _decode_flax_array(val)
|
|
# Slice first Dense layer: (254, 512) → (64, 512)
|
|
params['Dense_0'] = {
|
|
'kernel': params['Dense_0']['kernel'][:64, :],
|
|
'bias': params['Dense_0']['bias'],
|
|
}
|
|
return params
|
|
|
|
|
|
def load_state_preprocessor(path):
|
|
"""Load RunningStandardScaler stats (running_mean, running_variance) for the
|
|
first 64 dims = [code(19), obs(45)] that feed the actor.
|
|
|
|
CRITICAL: the policy was trained on NORMALIZED observations. Deployment must
|
|
apply: clip((x - mean) / (sqrt(var) + 1e-8), -5, 5) before the actor.
|
|
"""
|
|
with open(path, 'rb') as f:
|
|
ckpt = pickle.load(f)
|
|
if 'state_preprocessor' not in ckpt:
|
|
print("[WARN] No state_preprocessor in checkpoint — skipping normalization")
|
|
return None, None
|
|
sp = msgpack.unpackb(ckpt['state_preprocessor'], raw=False)
|
|
mean = _decode_flax_array(sp['running_mean'])[:64].astype(np.float32)
|
|
var = _decode_flax_array(sp['running_variance'])[:64].astype(np.float32)
|
|
return mean, var
|
|
|
|
|
|
def export_onnx(actor_params, vae_params, output_path, obs_mean=None, obs_var=None):
|
|
"""Build PyTorch model from Flax params, export to ONNX.
|
|
|
|
If obs_mean/obs_var given, bakes in the state-preprocessor normalization
|
|
(applied to [code(19), obs(45)] before the actor) — REQUIRED for the policy
|
|
to behave correctly, since it was trained on normalized observations.
|
|
"""
|
|
import torch, torch.nn as tnn
|
|
|
|
class DreamWaQTorch(tnn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
# CENet encoder (Manaro-Alpha: 225→128→64)
|
|
self.enc1 = tnn.Linear(225, 128)
|
|
self.enc2 = tnn.Linear(128, 64)
|
|
self.latent_mu = tnn.Linear(64, 16)
|
|
self.vel_mu = tnn.Linear(64, 3)
|
|
# Actor (64→512→256→128→12)
|
|
self.act1 = tnn.Linear(64, 512)
|
|
self.act2 = tnn.Linear(512, 256)
|
|
self.act3 = tnn.Linear(256, 128)
|
|
self.act_out = tnn.Linear(128, 12)
|
|
# State-preprocessor normalization buffers (for [code(19), obs(45)] = 64)
|
|
self.register_buffer("obs_mean", torch.zeros(64))
|
|
self.register_buffer("obs_std", torch.ones(64))
|
|
self.normalize = False
|
|
|
|
def forward(self, obs, history):
|
|
h = history.reshape(history.shape[0], -1)
|
|
h = tnn.functional.elu(self.enc1(h))
|
|
h = tnn.functional.elu(self.enc2(h))
|
|
z = self.latent_mu(h)
|
|
vel = self.vel_mu(h)
|
|
x = torch.cat([vel, z, obs], dim=-1)
|
|
# Apply state-preprocessor normalization (clip((x-mean)/(std+eps), -5, 5))
|
|
if self.normalize:
|
|
x = torch.clamp((x - self.obs_mean) / (self.obs_std + 1e-8), -5.0, 5.0)
|
|
x = tnn.functional.elu(self.act1(x))
|
|
x = tnn.functional.elu(self.act2(x))
|
|
x = tnn.functional.elu(self.act3(x))
|
|
return self.act_out(x)
|
|
|
|
model = DreamWaQTorch()
|
|
if obs_mean is not None and obs_var is not None:
|
|
model.obs_mean.data = torch.from_numpy(obs_mean.copy())
|
|
model.obs_std.data = torch.from_numpy(np.sqrt(obs_var).copy())
|
|
model.normalize = True
|
|
print("[ONNX] State-preprocessor normalization baked in")
|
|
|
|
# Transfer CENet encoder weights (from Flax frozen dict)
|
|
vp = vae_params['params']
|
|
model.enc1.weight.data = torch.from_numpy(np.array(vp['enc_fc1']['kernel']).T.copy())
|
|
model.enc1.bias.data = torch.from_numpy(np.array(vp['enc_fc1']['bias']).copy())
|
|
model.enc2.weight.data = torch.from_numpy(np.array(vp['enc_fc2']['kernel']).T.copy())
|
|
model.enc2.bias.data = torch.from_numpy(np.array(vp['enc_fc2']['bias']).copy())
|
|
model.latent_mu.weight.data = torch.from_numpy(np.array(vp['latent_mu']['kernel']).T.copy())
|
|
model.latent_mu.bias.data = torch.from_numpy(np.array(vp['latent_mu']['bias']).copy())
|
|
model.vel_mu.weight.data = torch.from_numpy(np.array(vp['vel_mu']['kernel']).T.copy())
|
|
model.vel_mu.bias.data = torch.from_numpy(np.array(vp['vel_mu']['bias']).copy())
|
|
|
|
# Transfer Actor weights (64-dim, already sliced, from msgpack decoded)
|
|
ap = actor_params
|
|
model.act1.weight.data = torch.from_numpy(ap['Dense_0']['kernel'].T.copy())
|
|
model.act1.bias.data = torch.from_numpy(ap['Dense_0']['bias'].copy())
|
|
model.act2.weight.data = torch.from_numpy(ap['Dense_1']['kernel'].T.copy())
|
|
model.act2.bias.data = torch.from_numpy(ap['Dense_1']['bias'].copy())
|
|
model.act3.weight.data = torch.from_numpy(ap['Dense_2']['kernel'].T.copy())
|
|
model.act3.bias.data = torch.from_numpy(ap['Dense_2']['bias'].copy())
|
|
model.act_out.weight.data = torch.from_numpy(ap['Dense_3']['kernel'].T.copy())
|
|
model.act_out.bias.data = torch.from_numpy(ap['Dense_3']['bias'].copy())
|
|
|
|
model.eval()
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
torch.onnx.export(
|
|
model,
|
|
(torch.randn(1, 45), torch.randn(1, 5, 45)),
|
|
output_path,
|
|
input_names=['observations', 'obs_history'],
|
|
output_names=['actions'],
|
|
opset_version=11,
|
|
dynamic_axes={'observations': {0: 'batch'}, 'obs_history': {0: 'batch'}, 'actions': {0: 'batch'}},
|
|
)
|
|
print(f"[ONNX] Exported → {output_path}")
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# Main
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--checkpoint", default=None, help="SKRL agent checkpoint")
|
|
p.add_argument("--vae", default=None, help="CENet params .pkl")
|
|
p.add_argument("--output", default=os.path.join(PROJECT, "exports_go1_dreamwaq", "policy.onnx"))
|
|
args = p.parse_args()
|
|
|
|
run_dir = os.path.join(PROJECT, "runs", "go1-dreamwaq-walk", "skrl")
|
|
|
|
# Auto-find checkpoint
|
|
if not args.checkpoint:
|
|
runs = sorted([d for d in os.listdir(run_dir) if os.path.isdir(os.path.join(run_dir, d)) and d.startswith("26-")])
|
|
if runs:
|
|
ckpt_dir = os.path.join(run_dir, runs[-1], "checkpoints")
|
|
args.checkpoint = os.path.join(ckpt_dir, "best_agent.pickle")
|
|
|
|
# Auto-find VAE params
|
|
if not args.vae:
|
|
vae_files = sorted([f for f in os.listdir(run_dir) if f.startswith("vae_") and f.endswith(".pkl")],
|
|
key=lambda x: int(x.split("_")[1].split(".")[0]))
|
|
if vae_files:
|
|
args.vae = os.path.join(run_dir, vae_files[-1])
|
|
else:
|
|
# Try cenet_params.pkl (saved at end of training)
|
|
cpath = os.path.join(run_dir, "cenet_params.pkl")
|
|
if os.path.exists(cpath):
|
|
args.vae = cpath
|
|
|
|
if not args.checkpoint or not os.path.exists(args.checkpoint):
|
|
print(f"[ERROR] Checkpoint not found: {args.checkpoint}")
|
|
sys.exit(1)
|
|
if not args.vae or not os.path.exists(args.vae):
|
|
print(f"[ERROR] VAE params not found: {args.vae}")
|
|
sys.exit(1)
|
|
|
|
print(f"Policy: {args.checkpoint}")
|
|
print(f"VAE: {args.vae}")
|
|
|
|
with open(args.vae, 'rb') as f:
|
|
vae_params = pickle.load(f)
|
|
actor_params = load_skrl_policy(args.checkpoint)
|
|
obs_mean, obs_var = load_state_preprocessor(args.checkpoint)
|
|
export_onnx(actor_params, vae_params, args.output, obs_mean, obs_var)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|