#!/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()