fix bug about collapse in mujoco

[HazinTihoa]

when i run the sim2tim cmd below:
ros2 launch motion_tracking_controller mujoco.launch.py policy_path:=~/whole_body_tracking/logs/rsl_rl/g1_flat/exported/policy.onnx
it reported error:

[mujoco_sim-3] [ERROR] [1762352259.045921687] [walking_controller]: Original error: OnnxPolicy: 'anchor_body_name' not found in model metadata. Please check the model.
[mujoco_sim-3] [WARN] [1762352259.045944157] [walking_controller]: Error occurred while doing error handling.
[mujoco_sim-3] [ERROR] [1762352259.045960838] [controller_manager]: After configuring, controller 'walking_controller' is in state 'unconfigured' , expected inactive.
[spawner-1] [ERROR] [1762352259.049161804] [spawner_state_estimator]: Failed to configure controller

basically, it due to attach_onnx_metadata() does not name anchor_body_name in .onnx file.

to fix this,you need to specify it in attach_onnx_matadata():

` def attach_onnx_metadata(env: ManagerBasedRLEnv, run_path: str, path: str, filename="policy.onnx") -> None:
onnx_path = os.path.join(path, filename)

motion_cmd: MotionCommand = env.command_manager.get_term("motion")                   # new

metadata = {
    "run_path": run_path,
    "joint_names": env.scene["robot"].data.joint_names,
    "joint_stiffness": env.scene["robot"].data.joint_stiffness[0].cpu().tolist(),
    "joint_damping": env.scene["robot"].data.joint_damping[0].cpu().tolist(),
    "default_joint_pos": env.scene["robot"].data.default_joint_pos_nominal.cpu().tolist(),
    "command_names": env.command_manager.active_terms,
    "observation_names": env.observation_manager.active_terms["policy"],
    "action_scale": env.action_manager.get_term("joint_pos")._scale[0].cpu().tolist(),
    "anchor_body_name": motion_cmd.cfg.anchor_body,                                                      #new
}

model = onnx.load(onnx_path)

for k, v in metadata.items():
    entry = onnx.StringStringEntryProto()
    entry.key = k
    entry.value = list_to_csv_str(v) if isinstance(v, list) else str(v)
    model.metadata_props.append(entry)

onnx.save(model, onnx_path)`

then run the play.py to export a new .onnx file

[fjc378209410]

import os
import torch

import onnx

from isaaclab.envs import ManagerBasedRLEnv
from isaaclab_rl.rsl_rl.exporter import _OnnxPolicyExporter

from whole_body_tracking.tasks.tracking.mdp import MotionCommand

def export_motion_policy_as_onnx(
env: ManagerBasedRLEnv,
actor_critic: object,
path: str,
normalizer: object | None = None,
filename="policy.onnx",
verbose=False,
):
if not os.path.exists(path):
os.makedirs(path, exist_ok=True)
policy_exporter = _OnnxMotionPolicyExporter(env, actor_critic, normalizer, verbose)
policy_exporter.export(path, filename)

class _OnnxMotionPolicyExporter(_OnnxPolicyExporter):
def init(self, env: ManagerBasedRLEnv, actor_critic, normalizer=None, verbose=False):
super().init(actor_critic, normalizer, verbose)
cmd: MotionCommand = env.command_manager.get_term("motion")

    self.joint_pos = cmd.motion.joint_pos.to("cpu")
    self.joint_vel = cmd.motion.joint_vel.to("cpu")
    self.body_pos_w = cmd.motion.body_pos_w.to("cpu")
    self.body_quat_w = cmd.motion.body_quat_w.to("cpu")
    self.body_lin_vel_w = cmd.motion.body_lin_vel_w.to("cpu")
    self.body_ang_vel_w = cmd.motion.body_ang_vel_w.to("cpu")
    self.time_step_total = self.joint_pos.shape[0]

def forward(self, x, time_step):
    time_step_clamped = torch.clamp(time_step.long().squeeze(-1), max=self.time_step_total - 1)
    return (
        self.actor(self.normalizer(x)),  # 这里会调用 normalizer 的 forward 方法
        self.joint_pos[time_step_clamped],
        self.joint_vel[time_step_clamped],
        self.body_pos_w[time_step_clamped],
        self.body_quat_w[time_step_clamped],
        self.body_lin_vel_w[time_step_clamped],
        self.body_ang_vel_w[time_step_clamped],
    )

def export(self, path, filename):
    self.to("cpu")
    # 关键改动1：激活归一化器的导出模式（关闭 update 中的动态判断）
    if self.normalizer is not None:
        self.normalizer.export_mode = True  # 对应之前在 normalizer.py 中添加的开关
    
    # 关键改动2：切换模型到评估模式（避免训练模式的动态行为）
    if hasattr(self.actor, "eval"):
        self.actor.eval()  # 切换 actor 到评估模式
    if self.normalizer is not None and hasattr(self.normalizer, "eval"):
        self.normalizer.eval()  # 切换 normalizer 到评估模式

    # 关键改动3：提高 opset 版本（日志提示旧版本可能不兼容，推荐用 18+）
    obs = torch.zeros(1, self.actor[0].in_features)
    time_step = torch.zeros(1, 1)
    torch.onnx.export(
        self,
        (obs, time_step),
        os.path.join(path, filename),
        export_params=True,
        opset_version=18,  # 从 11 改为 18，解决版本兼容警告
        verbose=self.verbose,
        input_names=["obs", "time_step"],
        output_names=[
            "actions",
            "joint_pos",
            "joint_vel",
            "body_pos_w",
            "body_quat_w",
            "body_lin_vel_w",
            "body_ang_vel_w",
        ],
        dynamic_axes={},  # 保持固定形状，避免动态维度导致的问题
    )

def list_to_csv_str(arr, *, decimals: int = 3, delimiter: str = ",") -> str:
fmt = f"{{:.{decimals}f}}"
return delimiter.join(
fmt.format(x) if isinstance(x, (int, float)) else str(x) for x in arr # numbers → format, strings → as-is
)

def attach_onnx_metadata(env: ManagerBasedRLEnv, run_path: str, path: str, filename="policy.onnx") -> None:
onnx_path = os.path.join(path, filename)

observation_names = env.observation_manager.active_terms["policy"]
observation_history_lengths: list[int] = []

if env.observation_manager.cfg.policy.history_length is not None:
    observation_history_lengths = [env.observation_manager.cfg.policy.history_length] * len(observation_names)
else:
    for name in observation_names:
        term_cfg = env.observation_manager.cfg.policy.to_dict()[name]
        history_length = term_cfg["history_length"]
        observation_history_lengths.append(1 if history_length == 0 else history_length)

metadata = {
    "run_path": run_path,
    "joint_names": env.scene["robot"].data.joint_names,
    "joint_stiffness": env.scene["robot"].data.joint_stiffness[0].cpu().tolist(),
    "joint_damping": env.scene["robot"].data.joint_damping[0].cpu().tolist(),
    "default_joint_pos": env.scene["robot"].data.default_joint_pos_nominal.cpu().tolist(),
    "command_names": env.command_manager.active_terms,
    "observation_names": observation_names,
    "observation_history_lengths": observation_history_lengths,
    "action_scale": env.action_manager.get_term("joint_pos")._scale[0].cpu().tolist(),
    "anchor_body_name": env.command_manager.get_term("motion").cfg.anchor_body_name,
    "body_names": env.command_manager.get_term("motion").cfg.body_names,
}

model = onnx.load(onnx_path)

for k, v in metadata.items():
    entry = onnx.StringStringEntryProto()
    entry.key = k
    entry.value = list_to_csv_str(v) if isinstance(v, list) else str(v)
    model.metadata_props.append(entry)

onnx.save(model, onnx_path)