98 lines
4.0 KiB
Python
98 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Export a DreamWaQ rsl_rl checkpoint (ActorCritic_DWAQ) to ONNX for MuJoCo.
|
|
|
|
rsl_rl has NO state preprocessor, so no normalization is needed (unlike SKRL).
|
|
The exported model uses the MEAN CENet code (deterministic deploy).
|
|
|
|
ONNX: inputs observations(1,45) + obs_history(1,5,45) -> actions(1,12)
|
|
(matches scripts/dreamwaq_sim2sim_mujoco.py interface)
|
|
|
|
Usage:
|
|
uv run scripts/export_dreamwaq_rsl_onnx.py # auto-find latest
|
|
uv run scripts/export_dreamwaq_rsl_onnx.py --checkpoint runs/.../model_700.pt
|
|
"""
|
|
import argparse, glob, os, sys
|
|
import torch
|
|
import torch.nn as tnn
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from motrix_rl.dwaq_rsl.actor_critic_dwaq import ActorCritic_DWAQ
|
|
|
|
PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
|
|
class DwaqInfer(tnn.Module):
|
|
"""Deterministic inference: CENet mean code + RunningStats normalize + actor."""
|
|
def __init__(self, ac: ActorCritic_DWAQ):
|
|
super().__init__()
|
|
self.encoder = ac.encoder
|
|
self.encode_mean_vel = ac.encode_mean_vel
|
|
self.encode_mean_latent = ac.encode_mean_latent
|
|
self.actor = ac.actor
|
|
# Bake RunningStats into the exported model
|
|
if hasattr(ac, 'actor_normalizer') and ac.actor_normalizer.count > 10:
|
|
self.register_buffer("norm_mean", ac.actor_normalizer.mean.clone())
|
|
self.register_buffer("norm_std", ac.actor_normalizer.var.sqrt().clone())
|
|
self.has_norm = True
|
|
else:
|
|
self.has_norm = False
|
|
|
|
def forward(self, obs, obs_history):
|
|
h = self.encoder(obs_history.reshape(obs_history.shape[0], -1)) # (B,225)->(B,64)
|
|
vel = self.encode_mean_vel(h) # (B,3) mean velocity estimate
|
|
latent = self.encode_mean_latent(h) # (B,16) mean latent
|
|
code = torch.cat([vel, latent], dim=-1) # (B,19) = [vel, latent]
|
|
x = torch.cat([code, obs], dim=-1) # (B,64) = [code, obs]
|
|
if self.has_norm:
|
|
x = torch.clamp((x - self.norm_mean) / (self.norm_std + 1e-8), -5.0, 5.0)
|
|
return self.actor(x) # (B,12)
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--checkpoint", default=None)
|
|
p.add_argument("--output", default=os.path.join(PROJECT, "exports_go1_dreamwaq", "policy.onnx"))
|
|
p.add_argument("--num-obs", type=int, default=45)
|
|
p.add_argument("--num-priv", type=int, default=235)
|
|
p.add_argument("--num-hist", type=int, default=5)
|
|
p.add_argument("--num-act", type=int, default=12)
|
|
args = p.parse_args()
|
|
|
|
if args.checkpoint is None:
|
|
runs = sorted(glob.glob(os.path.join(PROJECT, "runs", "go1-dreamwaq-walk", "rsl_dwaq", "*")),
|
|
key=os.path.getmtime)
|
|
if not runs:
|
|
print("[ERROR] no rsl_dwaq runs found"); sys.exit(1)
|
|
models = glob.glob(os.path.join(runs[-1], "model_*.pt"))
|
|
args.checkpoint = max(models, key=os.path.getmtime)
|
|
|
|
print(f"[rsl-ONNX] checkpoint: {args.checkpoint}")
|
|
cenet_out = 19
|
|
ac = ActorCritic_DWAQ(
|
|
args.num_obs + cenet_out, # actor in = 64
|
|
args.num_priv, # critic in = 235
|
|
args.num_act, # 12
|
|
args.num_hist * args.num_obs, # cenet in = 225
|
|
cenet_out, # 19
|
|
)
|
|
ckpt = torch.load(args.checkpoint, map_location="cpu")
|
|
ac.load_state_dict(ckpt["model_state_dict"])
|
|
ac.eval()
|
|
|
|
model = DwaqInfer(ac).eval()
|
|
os.makedirs(os.path.dirname(args.output), exist_ok=True)
|
|
torch.onnx.export(
|
|
model,
|
|
(torch.zeros(1, args.num_obs), torch.zeros(1, args.num_hist, args.num_obs)),
|
|
args.output,
|
|
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"[rsl-ONNX] exported -> {args.output} (no normalization; mean CENet code)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|