diff --git a/deploy/include/FSM/FSMState.h b/deploy/include/FSM/FSMState.h index cbbc2afde..c136a8726 100644 --- a/deploy/include/FSM/FSMState.h +++ b/deploy/include/FSM/FSMState.h @@ -56,6 +56,7 @@ class FSMState : public BaseState void pre_run() { lowstate->update(); + if(privilegedstate) privilegedstate->update(); if(keyboard) keyboard->update(); } @@ -66,5 +67,6 @@ class FSMState : public BaseState static std::unique_ptr lowcmd; static std::shared_ptr lowstate; + static std::shared_ptr privilegedstate; static std::shared_ptr keyboard; }; \ No newline at end of file diff --git a/deploy/robots/g1_29dof/config/config.yaml b/deploy/robots/g1_29dof/config/config.yaml index 8cc97fee1..bba3e3df9 100644 --- a/deploy/robots/g1_29dof/config/config.yaml +++ b/deploy/robots/g1_29dof/config/config.yaml @@ -14,6 +14,9 @@ FSM: Mimic_Gangnam_Style: id: 102 type: Mimic + PhysHSI: # PhysHSI 导入的 ONNX 策略 + id: 200 + type: RLBase Passive: transitions: @@ -61,13 +64,19 @@ FSM: ] ] Velocity: - transitions: + transitions: Passive: LT + B.on_pressed Mimic_Dance_102: LT(2s) + down.on_pressed Mimic_Gangnam_Style: LT(2s) + left.on_pressed + PhysHSI: LT(2s) + right.on_pressed policy_dir: config/policy/velocity # policy_dir: ../../../logs/rsl_rl/unitree_g1_29dof_velocity + PhysHSI: + transitions: + Passive: LT + B.on_pressed + Velocity: RB + X.on_pressed + policy_dir: config/policy/physhsi/v1/ Mimic_Dance_102: transitions: Passive: LT + B.on_pressed diff --git a/deploy/robots/g1_29dof/config/policy/physhsi/v1/exported/policy.onnx b/deploy/robots/g1_29dof/config/policy/physhsi/v1/exported/policy.onnx new file mode 100644 index 000000000..77e7b2b53 Binary files /dev/null and b/deploy/robots/g1_29dof/config/policy/physhsi/v1/exported/policy.onnx differ diff --git a/deploy/robots/g1_29dof/config/policy/physhsi/v1/params/deploy.yaml b/deploy/robots/g1_29dof/config/policy/physhsi/v1/params/deploy.yaml new file mode 100644 index 000000000..57ad2a3e4 --- /dev/null +++ b/deploy/robots/g1_29dof/config/policy/physhsi/v1/params/deploy.yaml @@ -0,0 +1,159 @@ +# ============================================================================= +# PhysHSI 多输入 ONNX 模型对应的 deploy.yaml +# ONNX 模型: policy_name_multi.onnx (7 个独立输入，由 wrap_onnx_multi_input.py 生成) +# +# ★ 数据流: C++ ObservationManager → obs_map → OrtRunner → ONNX 推理 → 动作 +# ★ 7 个 ONNX 输入 = 7 个 observation group, 每个 group 内 use_gym_history: true +# +# 738 维 = 6 时间步 × 每步 123 维 +# 每步 123 维组成 (严格按此顺序): +# [0-2] base_ang_vel × 0.25 (3) IMU 陀螺仪 +# [3-5] projected_gravity (3) IMU 重力投影 +# [6-34] joint_pos_rel (29) 关节位置 - 默认值 +# [35-63] joint_vel_rel × 0.05 (29) 关节速度 +# [64-78] end_effector_pos (15) 5个末端在本体坐标系位置 +# [79-107]last_action (29) 上一步动作 +# [108-122]task_obs (15) 任务感知 (箱子+目标点) +# +# use_gym_history: true → C++ 按时间步交织输出: +# [t-5_all123] [t-4_all123] [t-3_all123] [t-2_all123] [t-1_all123] [t_all123] +# ============================================================================= + +# ===== 机器人基础配置 ===== +joint_ids_map: [0, 6, 12, 1, 7, 13, 2, 8, 14, 3, 9, 15, 22, 4, 10, 16, 23, 5, 11, + 17, 24, 18, 25, 19, 26, 20, 27, 21, 28] +step_dt: 0.02 + +# ===== PD 控制参数 ===== +stiffness: [150, 150, 150, 300, 40, 40, + 150, 150, 150, 300, 40, 40, + 300, 300, 300, + 200, 200, 200, 100, 20, + 200, 200, 200, 100, 20] + +damping: [2, 2, 2, 4, 1, 1, + 2, 2, 2, 4, 1, 1, + 4, 4, 4, + 3, 3, 3, 1, 0.5, + 3, 3, 3, 1, 0.5] + +# ===== 默认关节角度 ===== +default_joint_pos: [-0.1, 0.0, 0.0, 0.3, -0.2, 0.0, + -0.1, 0.0, 0.0, 0.3, -0.2, 0.0, + 0.0, 0.0, 0.0, + 0.0, 0.1, 0.0, 1.2, 0.0, + 0.0, -0.1, 0.0, 1.2, 0.0] + +# ===== 命令 ===== +commands: + base_velocity: + ranges: + lin_vel_x: [0.0, 1.0] + lin_vel_y: [-0.3, 0.3] + ang_vel_z: [-0.3, 0.3] + heading: null + +# ===== 动作 ===== +actions: + JointPositionAction: + clip: null + joint_names: [.*] + scale: 0.25 + # offset 自动使用 default_joint_pos + +# ============================================================================= +# 观测 — 多 Group 模式: 7 个 group → obs_map 7 个 key → 匹配 7 个 ONNX 输入 +# +# C++ ObservationManager._prapare_terms() 检测: +# - 第一个 key 的值无 "params" 键 → multi-group 模式 +# - 每个 group 名 = obs_map key = ONNX 输入名 (必须严格一致) +# - _prepare_group_terms 处理 group 内的 use_gym_history 和 observation terms +# +# C++ compute_group("base_ang_vel") 输出 (use_gym_history=true, history_length=6): +# h=0: term.get(0) = t-5 的 base_ang_vel [3] +# h=1: term.get(1) = t-4 的 base_ang_vel [3] +# ... ... +# h=5: term.get(5) = t 的 base_ang_vel [3] +# → [18] +# ============================================================================= +observations: + + # ── Group 1: base_ang_vel → ONNX 输入 [B, 18] ── + # 3 dims × 6 history = 18, scale = 0.25 + base_ang_vel: + use_gym_history: true + base_ang_vel: + params: {} + scale: [0.25, 0.25, 0.25] + history_length: 6 + + # ── Group 2: projected_gravity → ONNX 输入 [B, 18] ── + # 3 dims × 6 history = 18, scale = 1.0 + projected_gravity: + use_gym_history: true + projected_gravity: + params: {} + scale: [1.0, 1.0, 1.0] + history_length: 6 + + # ── Group 3: joint_pos_rel → ONNX 输入 [B, 174] ── + # 29 dims × 6 history = 174, scale = 1.0 + joint_pos_rel: + use_gym_history: true + joint_pos_rel: + params: {} + scale: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + history_length: 6 + + # ── Group 4: joint_vel_rel → ONNX 输入 [B, 174] ── + # 29 dims × 6 history = 174, scale = 0.05 + joint_vel_rel: + use_gym_history: true + joint_vel_rel: + params: {} + scale: [0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, + 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, + 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05] + history_length: 6 + + # ── Group 5: end_effector_pos → ONNX 输入 [B, 90] ── + # 15 dims × 6 history = 90, scale = 1.0 + # ★ 组成: left_palm(3) | right_palm(3) | left_ankle(3) | right_ankle(3) | d455_head(3) + # ★ 计算: quat_rotate_inverse(root_quat, pos_global - root_pos_global) + # ⚠ 需要在 C++ 侧注册 REGISTER_OBSERVATION(end_effector_pos) 函数 + end_effector_pos: + use_gym_history: true + end_effector_pos: + params: {} + scale: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0] + history_length: 6 + + # ── Group 6: last_action → ONNX 输入 [B, 174] ── + # 29 dims × 6 history = 174, scale = 1.0 + last_action: + use_gym_history: true + last_action: + params: {} + scale: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + history_length: 6 + + # ── Group 7: task_obs → ONNX 输入 [B, 90] ── + # 15 dims × 6 history = 90, scale = 1.0 + # ★ 组成: box_pos_local(3) | box_rot_6d_local(6) | box_size(3) | goal_pos_local(3) + # ★ box_pos_local: quat_rotate_inverse(root_quat, box_pos - root_pos) + # ★ box_rot_6d_local: quat_to_tan_norm(conjugate(root_quat) * box_quat) + # ★ box_size: 静态 [width, depth, height] + # ★ goal_pos_local: quat_rotate_inverse(root_quat, goal_pos - root_pos) + # ⚠ 需要在 C++ 侧注册 REGISTER_OBSERVATION(task_obs) 函数 + task_obs: + use_gym_history: true + task_obs: + params: {} + scale: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0] + history_length: 6 diff --git a/deploy/robots/g1_29dof/idl/PrivilegedState.idl b/deploy/robots/g1_29dof/idl/PrivilegedState.idl new file mode 100644 index 000000000..c40f54bd6 --- /dev/null +++ b/deploy/robots/g1_29dof/idl/PrivilegedState.idl @@ -0,0 +1,46 @@ +// ============================================================================= +// PrivilegedState.idl — 特权信息 DDS 消息类型定义 +// +// 放置路径: unitree_sdk2 的 unitree/idl/hg/ 目录下 +// 或在本项目的 deploy/robots/g1_29dof/idl/ 下 (需配置 CMake 的 IDL 编译) +// +// 编译命令 (Cyclone DDS): +// idlc -l cpp PrivilegedState.idl +// +// 生成文件: PrivilegedState_.hpp (C++ 头文件, 含 DDS 类型支持) +// 命名空间: unitree_hg::msg::dds_::PrivilegedState_ +// ============================================================================= + +module unitree_hg { + module msg { + + @topic + struct PrivilegedState_ + { + // 末端执行器世界坐标 (5 bodies × 3 = 15 floats) + float end_effector_pos_w[15]; + + // 机器人根Link世界坐标 (3 floats) + float root_pos_w[3]; + + // 箱子世界坐标 (3 floats) + float box_pos_w[3]; + + // 箱子世界姿态四元数 w,x,y,z (4 floats) + float box_quat_w[4]; + + // 箱子半边长 (3 floats, 静态) + float box_size[3]; + + // 目标点世界坐标 (3 floats) + float goal_pos_w[3]; + + // 帧序号 + unsigned long sequence; + + // 对齐填充 + unsigned long padding; + }; + + }; +}; diff --git a/deploy/robots/g1_29dof/include/PrivilegedStateSub.h b/deploy/robots/g1_29dof/include/PrivilegedStateSub.h new file mode 100644 index 000000000..db6a5879a --- /dev/null +++ b/deploy/robots/g1_29dof/include/PrivilegedStateSub.h @@ -0,0 +1,78 @@ +// Copyright (c) 2025, Unitree Robotics Co., Ltd. +// All rights reserved. +// +// 特权信息 DDS 订阅器 — 完全遵循 ros2_sub.h 参考模式 +// +// 模式: 继承 SubscriptionBase, 传入 topic 名称 +// 参考: CameraData : SubscriptionBase +// TorsoImu : SubscriptionBase +// 新增: PrivilegedState : SubscriptionBase +// +// 消息类型 PrivilegedState_ 的 IDL 定义见下方注释; +// 需在 unitree_sdk2 的 unitree/idl/hg/ 目录下添加 IDL 文件并用 idlc 生成 .hpp +// +// IDL 定义 (unitree/idl/hg/PrivilegedState.idl): +// +// module unitree_hg { +// module msg { +// struct PrivilegedState_ { +// float end_effector_pos_w[15]; // 5 bodies × 3 +// float root_pos_w[3]; +// float box_pos_w[3]; +// float box_quat_w[4]; // w, x, y, z +// float box_size[3]; +// float goal_pos_w[3]; +// unsigned long sequence; +// unsigned long padding; +// }; +// }; +// }; + +#ifndef PRIVILEGED_STATE_SUB_H +#define PRIVILEGED_STATE_SUB_H + +#include +#include "unitree/dds_wrapper/common/Subscription.h" +#include "unitree/dds_wrapper/robots/g1/defines.h" + +// ★ IDL 生成的消息类型头文件 (需要在 unitree_sdk2 中添加 IDL 并编译生成) +#include "PrivilegedState_.hpp" + +namespace unitree +{ +namespace robot +{ +namespace g1 +{ +namespace subscription +{ + +// ============================================================================= +// PrivilegedState — 特权信息 DDS 订阅器 +// +// 完全遵循 CameraData / TorsoImu 的模式: +// - 继承 SubscriptionBase (unitree_sdk2 的 DDS 订阅封装) +// - 构造函数传入 topic 名称 +// - 自动获得 update() / wait_for_connection() / isTimeout() / msg_ / mutex_ +// +// 话题: rt/privileged_state +// 类型: unitree_hg::msg::dds_::PrivilegedState_ +// 域: 与 ChannelFactory::Init() 指定的同一 DDS 域 +// ============================================================================= + +class PrivilegedState : public SubscriptionBase +{ +public: + using SharedPtr = std::shared_ptr; + + PrivilegedState(std::string topic = "rt/privileged_state") + : SubscriptionBase(topic) + {} +}; + +} // namespace subscription +} // namespace g1 +} // namespace robot +} // namespace unitree + +#endif // PRIVILEGED_STATE_SUB_H diff --git a/deploy/robots/g1_29dof/include/PrivilegedState_.hpp b/deploy/robots/g1_29dof/include/PrivilegedState_.hpp new file mode 100644 index 000000000..fa9736356 --- /dev/null +++ b/deploy/robots/g1_29dof/include/PrivilegedState_.hpp @@ -0,0 +1,118 @@ +// ============================================================================= +// PrivilegedState_.hpp — 特权信息 DDS 消息类型定义 (替 idlc 生成) +// +// 来源: idl/PrivilegedState.idl +// 用途: SubscriptionBase 的模板参数 +// 经 Cyclone DDS C++11 API 自动序列化，无需 idlc +// +// 原理: Cyclone DDS C++11 对 POD struct 可通过 topic_type_support 特化 +// 实现零拷贝序列化 (128 bytes 直传，无动态字段) +// ============================================================================= + +#ifndef UNITREE_HG_MSG_DDS__PRIVILEGED_STATE__HPP +#define UNITREE_HG_MSG_DDS__PRIVILEGED_STATE__HPP + +#include +#include + +namespace unitree_hg { +namespace msg { +namespace dds_ { + +// ============================================================================= +// PrivilegedState_ — 128 bytes POD, DDS 兼容 +// +// ★ 字段顺序、类型、数组大小必须与 IDL 定义严格一致 +// Cyclone DDS CDR 序列化依赖 struct 内存布局 +// ============================================================================= + +class PrivilegedState_ +{ +public: + // ---- 原始数据访问 (符合 IDL 生成代码的访问器命名惯例) ---- + float* end_effector_pos_w() { return m_end_effector_pos_w; } + const float* end_effector_pos_w() const { return m_end_effector_pos_w; } + + float* root_pos_w() { return m_root_pos_w; } + const float* root_pos_w() const { return m_root_pos_w; } + + float* box_pos_w() { return m_box_pos_w; } + const float* box_pos_w() const { return m_box_pos_w; } + + float* box_quat_w() { return m_box_quat_w; } + const float* box_quat_w() const { return m_box_quat_w; } + + float* box_size() { return m_box_size; } + const float* box_size() const { return m_box_size; } + + float* goal_pos_w() { return m_goal_pos_w; } + const float* goal_pos_w() const { return m_goal_pos_w; } + + uint32_t& sequence() { return m_sequence; } + const uint32_t& sequence() const { return m_sequence; } + + uint32_t& padding() { return m_padding; } + const uint32_t& padding() const { return m_padding; } + + // DDS 要求: 默认构造函数 zero-initialize + PrivilegedState_() { std::memset(this, 0, sizeof(PrivilegedState_)); } + +private: + float m_end_effector_pos_w[15]; // offset 0, 60 bytes + float m_root_pos_w[3]; // offset 60, 12 bytes + float m_box_pos_w[3]; // offset 72, 12 bytes + float m_box_quat_w[4]; // offset 84, 16 bytes + float m_box_size[3]; // offset 100,12 bytes + float m_goal_pos_w[3]; // offset 112,12 bytes + uint32_t m_sequence; // offset 124,4 bytes + uint32_t m_padding; // offset 128,4 bytes → total 128 +}; + +// 编译期验证大小 +static_assert(sizeof(PrivilegedState_) == 128, + "PrivilegedState_ must be exactly 128 bytes"); +static_assert(alignof(PrivilegedState_) <= 8, + "PrivilegedState_ alignment must not exceed 8 bytes"); + +} // namespace dds_ +} // namespace msg +} // namespace unitree_hg + +// ============================================================================= +// DDS 类型注册 — 告诉 Cyclone DDS 这是合法的 topic 类型 +// +// Cyclone DDS C++11 通过 dds::topic::is_topic 和 +// dds::topic::topic_type_support 两个特化来识别 topic 类型。 +// +// 对于纯 POD struct (无指针、无字符串、固定大小), +// Cyclone DDS 使用默认 CDR 序列化: 直接 memcpy 128 bytes, +// 无需手写 serialize/deserialize 函数。 +// ============================================================================= + +#include + +namespace dds { +namespace topic { + +template<> +struct is_topic +{ + static constexpr bool value = true; +}; + +template<> +struct topic_type_support +{ + // 对于固定大小的 POD 类型，返回类型名用于 DDS 发现匹配 + static const std::string& name() + { + static const std::string n = + "unitree_hg::msg::dds_::PrivilegedState_"; + return n; + } +}; + +} // namespace topic +} // namespace dds + +#endif // UNITREE_HG_MSG_DDS__PRIVILEGED_STATE__HPP diff --git a/deploy/robots/g1_29dof/include/Types.h b/deploy/robots/g1_29dof/include/Types.h index ee873a59e..c535cb7b9 100644 --- a/deploy/robots/g1_29dof/include/Types.h +++ b/deploy/robots/g1_29dof/include/Types.h @@ -3,5 +3,8 @@ #include "unitree/dds_wrapper/robots/go2/go2.h" #include "unitree/dds_wrapper/robots/g1/g1.h" +#include "PrivilegedStateSub.h" + using LowCmd_t = unitree::robot::g1::publisher::LowCmd; -using LowState_t = unitree::robot::g1::subscription::LowState; \ No newline at end of file +using LowState_t = unitree::robot::g1::subscription::LowState; +using PrivilegedState_t = unitree::robot::g1::subscription::PrivilegedState; \ No newline at end of file diff --git a/deploy/robots/g1_29dof/main.cpp b/deploy/robots/g1_29dof/main.cpp index be4041872..3556b9a4a 100644 --- a/deploy/robots/g1_29dof/main.cpp +++ b/deploy/robots/g1_29dof/main.cpp @@ -6,6 +6,7 @@ std::unique_ptr FSMState::lowcmd = nullptr; std::shared_ptr FSMState::lowstate = nullptr; +std::shared_ptr FSMState::privilegedstate = nullptr; std::shared_ptr FSMState::keyboard = std::make_shared(); void init_fsm_state() @@ -23,6 +24,12 @@ void init_fsm_state() spdlog::info("Waiting for connection to robot..."); FSMState::lowstate->wait_for_connection(); spdlog::info("Connected to robot."); + + // 订阅特权信息话题 (遵循 LowState 的 SubscriptionBase 模式) + FSMState::privilegedstate = std::make_shared(); + spdlog::info("Waiting for privileged state publisher..."); + FSMState::privilegedstate->wait_for_connection(); + spdlog::info("Privileged state topic connected."); } int main(int argc, char** argv) @@ -43,7 +50,7 @@ int main(int argc, char** argv) spdlog::critical("Unmatched robot type."); exit(-1); } - + // Initialize FSM auto fsm = std::make_unique(param::config["FSM"]); fsm->start(); @@ -55,7 +62,6 @@ int main(int argc, char** argv) { sleep(1); } - + return 0; } - diff --git a/deploy/robots/g1_29dof/src/PhysHSI_Observations.cpp b/deploy/robots/g1_29dof/src/PhysHSI_Observations.cpp new file mode 100644 index 000000000..1f6a525c1 --- /dev/null +++ b/deploy/robots/g1_29dof/src/PhysHSI_Observations.cpp @@ -0,0 +1,160 @@ +// Copyright (c) 2025, Unitree Robotics Co., Ltd. +// All rights reserved. +// +// PhysHSI 特有 Observation 函数注册 +// +// 遵循参考 commit 的模式: +// CameraData → REGISTER_OBSERVATION(depth_image) → asset->data.depth_image_buffer +// TorsoImu → (通过 CameraArticulation 访问) +// PrivilegedState → REGISTER_OBSERVATION(end_effector_pos / task_obs) +// → FSMState::privilegedstate->msg_ (与 lowstate->msg_ 同模式) +// +// 数据来源: unitree_mujoco 通过 DDS topic rt/privileged_state 发布 +// 订阅方式: PrivilegedState_t (继承 SubscriptionBase, 与 LowState_t 相同模式) + +#include "FSM/FSMState.h" +#include "unitree_articulation.h" +#include "isaaclab/envs/mdp/observations/observations.h" + +namespace isaaclab +{ +namespace mdp +{ + +// ============================================================================= +// 辅助函数 +// ============================================================================= + +// quat_rotate_inverse(q, v_world) — 世界坐标转本体坐标 +inline Eigen::Vector3f quat_rotate_inverse(const Eigen::Quaternionf& q, + const Eigen::Vector3f& v) +{ + return q.conjugate() * v; +} + +// quat_to_tan_norm(q) — 四元数转 6D 旋转表示 +// 对应 Python: mat = matrix_from_quat(q); mat[:, :2].reshape(-1) +// 输出顺序: [r00, r01, r10, r11, r20, r21] (行主序, 与 PyTorch 一致) +inline Eigen::Matrix quat_to_tan_norm(const Eigen::Quaternionf& q) +{ + Eigen::Matrix3f R = q.normalized().toRotationMatrix(); + Eigen::Matrix result; + result << R(0, 0), R(0, 1), R(1, 0), R(1, 1), R(2, 0), R(2, 1); + return result; +} + + +// ============================================================================= +// REGISTER_OBSERVATION: end_effector_pos (15 维) +// +// 访问模式: FSMState::privilegedstate->msg_ (与 lowstate->msg_ 同模式) +// ============================================================================= + +REGISTER_OBSERVATION(end_effector_pos) +{ + auto& robot = env->robot; + auto root_quat = robot->data.root_quat_w; + + std::vector result(15, 0.0f); + + if (FSMState::privilegedstate) + { + std::lock_guard lock(FSMState::privilegedstate->mutex_); + auto& msg = FSMState::privilegedstate->msg_; + + // 从 DDS 消息中构建 root_pos_w + Eigen::Vector3f root_pos_w(msg.root_pos_w()[0], + msg.root_pos_w()[1], + msg.root_pos_w()[2]); + + // 5 个末端在本体坐标系中的位置 + for (int i = 0; i < 5; ++i) + { + Eigen::Vector3f pos_w(msg.end_effector_pos_w()[i * 3 + 0], + msg.end_effector_pos_w()[i * 3 + 1], + msg.end_effector_pos_w()[i * 3 + 2]); + + Eigen::Vector3f pos_b = quat_rotate_inverse(root_quat, + pos_w - root_pos_w); + + result[i * 3 + 0] = pos_b.x(); + result[i * 3 + 1] = pos_b.y(); + result[i * 3 + 2] = pos_b.z(); + } + } + + return result; +} + + +// ============================================================================= +// REGISTER_OBSERVATION: task_obs (15 维) +// +// 访问模式: FSMState::privilegedstate->msg_ (与 lowstate->msg_ 同模式) +// ============================================================================= + +REGISTER_OBSERVATION(task_obs) +{ + auto& robot = env->robot; + auto root_quat = robot->data.root_quat_w; + + std::vector result(15, 0.0f); + + if (FSMState::privilegedstate) + { + std::lock_guard lock(FSMState::privilegedstate->mutex_); + auto& msg = FSMState::privilegedstate->msg_; + + Eigen::Vector3f root_pos_w(msg.root_pos_w()[0], + msg.root_pos_w()[1], + msg.root_pos_w()[2]); + + // ---- (A) box_pos_local (3维) ---- + { + Eigen::Vector3f box_w(msg.box_pos_w()[0], + msg.box_pos_w()[1], + msg.box_pos_w()[2]); + Eigen::Vector3f box_b = quat_rotate_inverse(root_quat, + box_w - root_pos_w); + result[0] = box_b.x(); + result[1] = box_b.y(); + result[2] = box_b.z(); + } + + // ---- (B) box_rot_6d_local (6维) ---- + { + Eigen::Quaternionf box_quat_w(msg.box_quat_w()[0], // w + msg.box_quat_w()[1], // x + msg.box_quat_w()[2], // y + msg.box_quat_w()[3]); // z + Eigen::Quaternionf box_quat_b = root_quat.conjugate() * box_quat_w; + Eigen::Matrix r6d = quat_to_tan_norm(box_quat_b); + for (int i = 0; i < 6; ++i) + result[3 + i] = r6d(i); + } + + // ---- (C) box_size (3维, 静态) ---- + { + result[9] = msg.box_size()[0]; + result[10] = msg.box_size()[1]; + result[11] = msg.box_size()[2]; + } + + // ---- (D) goal_pos_local (3维) ---- + { + Eigen::Vector3f goal_w(msg.goal_pos_w()[0], + msg.goal_pos_w()[1], + msg.goal_pos_w()[2]); + Eigen::Vector3f goal_b = quat_rotate_inverse(root_quat, + goal_w - root_pos_w); + result[12] = goal_b.x(); + result[13] = goal_b.y(); + result[14] = goal_b.z(); + } + } + + return result; +} + +} // namespace mdp +} // namespace isaaclab diff --git a/doc/unitree_mujoco_privileged_publisher_ref.cpp b/doc/unitree_mujoco_privileged_publisher_ref.cpp new file mode 100644 index 000000000..a17c109ba --- /dev/null +++ b/doc/unitree_mujoco_privileged_publisher_ref.cpp @@ -0,0 +1,227 @@ +// ============================================================================= +// unitree_mujoco 侧 — 特权信息 DDS Publisher 参考实现 +// +// 话题: rt/privileged_state (Cyclone DDS, 域 0) +// 类型: PrivilegedStateMsg (128 bytes POD) +// +// 与 g1_ctrl 的 PrivilegedSubscriber 配套使用。 +// 双方在同一个 DDS 域自动发现、自动匹配。 +// +// 集成方式: 将此代码合并到 unitree_mujoco/simulate/ 仿真主循环中。 +// ============================================================================= + +#include // Cyclone DDS C++11 API +#include // MJCF (根据实际路径调整) +#include +#include + +// ============================================================================= +// PrivilegedStateMsg — 与 g1_ctrl 侧 struct 完全一致 (128 bytes) +// ============================================================================= + +struct PrivilegedStateMsg +{ + float end_effector_pos_w[15]; // 5 bodies × 3 (60 bytes) + float root_pos_w[3]; // root link pos (12 bytes) + float box_pos_w[3]; // box pos (12 bytes) + float box_quat_w[4]; // box quat w,x,y,z (16 bytes) + float box_size[3]; // box half-size (12 bytes) + float goal_pos_w[3]; // goal pos (12 bytes) + uint32_t sequence; // frame counter (4 bytes) + uint32_t padding; // alignment (4 bytes) +}; +static_assert(sizeof(PrivilegedStateMsg) == 128, "Must be 128 bytes"); + + +// ============================================================================= +// 全局 DDS 状态 +// ============================================================================= + +static std::unique_ptr g_dp; +static std::unique_ptr> g_topic; +static std::unique_ptr g_pub; +static std::unique_ptr> g_writer; +static uint32_t g_sequence = 0; + +// 缓存的 body/site ID +static int g_body_left_palm = -1; +static int g_body_right_palm = -1; +static int g_body_left_ankle = -1; +static int g_body_right_ankle = -1; +static int g_body_head = -1; +static int g_body_root = -1; +static int g_body_box = -1; +static int g_site_goal = -1; +static int g_geom_box = -1; + +static constexpr const char* TOPIC_NAME = "rt/privileged_state"; + + +// ============================================================================= +// init_privileged_publisher — 初始化 DDS DataWriter (仿真启动时调用一次) +// ============================================================================= + +bool init_privileged_publisher(mjModel* m) +{ + // ---- 1. 缓存 MuJoCo body/site/geom ID ---- + g_body_left_palm = mj_name2id(m, mjOBJ_BODY, "left_palm_link"); + g_body_right_palm = mj_name2id(m, mjOBJ_BODY, "right_palm_link"); + g_body_left_ankle = mj_name2id(m, mjOBJ_BODY, "left_ankle_pitch_link"); + g_body_right_ankle = mj_name2id(m, mjOBJ_BODY, "right_ankle_pitch_link"); + g_body_head = mj_name2id(m, mjOBJ_BODY, "d455_link"); + g_body_root = mj_name2id(m, mjOBJ_BODY, "torso_link"); + g_body_box = mj_name2id(m, mjOBJ_BODY, "box"); + g_site_goal = mj_name2id(m, mjOBJ_SITE, "goal"); + + if (g_body_box >= 0) + { + int geom_start = m->body_geomadr[g_body_box]; + int geom_num = m->body_geomnum[g_body_box]; + if (geom_num > 0 && m->geom_type[geom_start] == mjGEOM_BOX) + g_geom_box = geom_start; + } + + // ---- 2. 创建 DDS Participant (域 0, 与 unitree_sdk2 同一域) ---- + try + { + g_dp = std::make_unique(0); + + // 创建话题 (Cyclone DDS 自动处理 POD 类型) + g_topic = std::make_unique>( + *g_dp, TOPIC_NAME); + + // 创建 Publisher + g_pub = std::make_unique(*g_dp); + + // DataWriter QoS: Reliable + KeepLast(10) + dds::pub::qos::DataWriterQos dw_qos; + dw_qos << dds::core::policy::Reliability::Reliable(); + dw_qos << dds::core::policy::Durability::TransientLocal(); + dw_qos << dds::core::policy::History::KeepLast(10); + + g_writer = std::make_unique>( + *g_pub, *g_topic, dw_qos); + + printf("[PrivilegedPublisher] DDS topic '%s' ready (domain 0)\n", + TOPIC_NAME); + return true; + } + catch (const dds::core::Exception& e) + { + printf("[PrivilegedPublisher] DDS init failed: %s\n", e.what()); + return false; + } +} + + +// ============================================================================= +// publish_privileged_state — 发布一帧特权信息 (每个仿真步调用) +// ============================================================================= + +void publish_privileged_state(mjModel* m, mjData* d) +{ + if (!g_writer) return; + + PrivilegedStateMsg msg; + std::memset(&msg, 0, sizeof(msg)); + msg.sequence = g_sequence++; + + // ---- (A) 末端执行器世界坐标 ---- + auto copy_body = [&](int id, float* dst) { + if (id >= 0) { + dst[0] = d->xpos[id*3+0]; + dst[1] = d->xpos[id*3+1]; + dst[2] = d->xpos[id*3+2]; + } + }; + copy_body(g_body_left_palm, msg.end_effector_pos_w + 0); + copy_body(g_body_right_palm, msg.end_effector_pos_w + 3); + copy_body(g_body_left_ankle, msg.end_effector_pos_w + 6); + copy_body(g_body_right_ankle, msg.end_effector_pos_w + 9); + copy_body(g_body_head, msg.end_effector_pos_w + 12); + + // ---- (B) 根Link位置 ---- + copy_body(g_body_root, msg.root_pos_w); + + // ---- (C) 箱子位置和姿态 ---- + if (g_body_box >= 0) + { + msg.box_pos_w[0] = d->xpos[g_body_box*3+0]; + msg.box_pos_w[1] = d->xpos[g_body_box*3+1]; + msg.box_pos_w[2] = d->xpos[g_body_box*3+2]; + + // MuJoCo xquat: [w, x, y, z] + msg.box_quat_w[0] = d->xquat[g_body_box*4+0]; + msg.box_quat_w[1] = d->xquat[g_body_box*4+1]; + msg.box_quat_w[2] = d->xquat[g_body_box*4+2]; + msg.box_quat_w[3] = d->xquat[g_body_box*4+3]; + } + + // ---- (D) 箱子尺寸 (从 mjModel 读取, 静态) ---- + if (g_geom_box >= 0) + { + msg.box_size[0] = m->geom_size[g_geom_box*3+0]; + msg.box_size[1] = m->geom_size[g_geom_box*3+1]; + msg.box_size[2] = m->geom_size[g_geom_box*3+2]; + } + + // ---- (E) 目标点 ---- + if (g_site_goal >= 0) + { + msg.goal_pos_w[0] = d->site_xpos[g_site_goal*3+0]; + msg.goal_pos_w[1] = d->site_xpos[g_site_goal*3+1]; + msg.goal_pos_w[2] = d->site_xpos[g_site_goal*3+2]; + } + + // ---- (F) 发布到 DDS 话题 ---- + try + { + g_writer->write(msg); + } + catch (const dds::core::Exception& e) + { + // 写入失败 (如 QoS 不匹配), 非致命 + } +} + + +// ============================================================================= +// cleanup_privileged_publisher — 清理 DDS 资源 (仿真退出时调用) +// ============================================================================= + +void cleanup_privileged_publisher() +{ + g_writer.reset(); + g_pub.reset(); + g_topic.reset(); + g_dp.reset(); + printf("[PrivilegedPublisher] DDS resources released\n"); +} + + +// ============================================================================= +// 集成示例 — unitree_mujoco 仿真主循环: +// +// int main() { +// mjModel* m = mj_loadXML("scene.xml", nullptr, ...); +// mjData* d = mj_makeData(m); +// +// // ★ 初始化 DDS 通信 (unitree_sdk2 已在别处调用 ChannelFactory::Init) +// init_privileged_publisher(m); +// +// while (running) { +// mj_step(m, d); +// +// // ... 现有: 发布 DDS LowState ... +// +// // ★ 发布特权信息 (与 LowState 在同一 DDS 总线) +// publish_privileged_state(m, d); +// +// // ... 现有: 订阅 DDS LowCmd, 写入 mjData->ctrl ... +// } +// +// cleanup_privileged_publisher(); +// mj_deleteData(d); +// mj_deleteModel(m); +// } +// ============================================================================= diff --git a/physhsi_deploy.yaml b/physhsi_deploy.yaml new file mode 100644 index 000000000..acb1fee7b --- /dev/null +++ b/physhsi_deploy.yaml @@ -0,0 +1,88 @@ +# ===== 机器人基础配置 ===== +joint_ids_map: [0, 6, 12, 1, 7, 13, 2, 8, 14, 3, 9, 15, 22, 4, 10, 16, 23, 5, 11, + 17, 24, 18, 25, 19, 26, 20, 27, 21, 28] +step_dt: 0.02 # sim_dt(0.005) × decimation(4) + +# ===== PD 控制参数 (carrybox 与 standup 相同) ===== +stiffness: [150, 150, 150, 300, 40, 40, + 150, 150, 150, 300, 40, 40, + 300, 300, 300, + 200, 200, 200, 100, 20, + 200, 200, 200, 100, 20] + +damping: [2, 2, 2, 4, 1, 1, + 2, 2, 2, 4, 1, 1, + 4, 4, 4, + 3, 3, 3, 1, 0.5, + 3, 3, 3, 1, 0.5] + +# ===== 默认关节角度 (carrybox 使用站立位姿) ===== +default_joint_pos: [-0.1, 0.0, 0.0, 0.3, -0.2, 0.0, + -0.1, 0.0, 0.0, 0.3, -0.2, 0.0, + 0.0, 0.0, 0.0, + 0.0, 0.1, 0.0, 1.2, 0.0, + 0.0, -0.1, 0.0, 1.2, 0.0] + +# ===== 命令 ===== +commands: + base_velocity: + ranges: + lin_vel_x: [0.0, 1.0] + lin_vel_y: [-0.3, 0.3] + ang_vel_z: [-0.3, 0.3] + heading: null + +# ===== 动作 ===== +actions: + JointPositionAction: + clip: null + joint_names: [.*] + scale: 0.25 + # offset 自动使用 default_joint_pos + +# ===== 观测 (★ 核心：use_gym_history: true + history_length: 6) ===== +observations: + use_gym_history: true # ★ 按时间步交织 (与 carrybox obs_buf 格式一致) + + base_ang_vel: + params: {} + scale: [0.25, 0.25, 0.25] + history_length: 6 + + projected_gravity: + params: {} + scale: [1.0, 1.0, 1.0] + history_length: 6 + + joint_pos_rel: + params: {} + scale: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + history_length: 6 + + joint_vel_rel: + params: {} + scale: [0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, + 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, + 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05] + history_length: 6 + + end_effector_pos: # ★ PhysHSI 特有 + params: {} + scale: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0] + history_length: 6 + + last_action: + params: {} + scale: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + history_length: 6 + + task_obs: # ★ PhysHSI 特有 (箱子+目标点) + params: {} + scale: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, + 1.0, 1.0, 1.0, 1.0, 1.0] + history_length: 6 \ No newline at end of file diff --git a/policy_name.onnx b/policy_name.onnx new file mode 100755 index 000000000..a79617dc6 Binary files /dev/null and b/policy_name.onnx differ diff --git a/wrap_onnx_multi_input.py b/wrap_onnx_multi_input.py new file mode 100644 index 000000000..a2e25f9a1 --- /dev/null +++ b/wrap_onnx_multi_input.py @@ -0,0 +1,316 @@ +#!/usr/bin/env python3 +""" +ONNX 模型包装脚本：将单输入 ONNX (input [B, 738] → output [B, 29]) +包装为多输入 ONNX，每个 observation term 作为独立输入。 + +★ 关键: use_gym_history: true 时，C++ ObservationManager 的输出是 + 逐时间步交织 (per-time-step interleaving)，而非逐观测项分组。 + + 原始模型输入顺序: + [t-5: all_terms(123)] [t-4: all_terms(123)] ... [t: all_terms(123)] = 738 + + 包装逻辑: + 7 个 per-term 输入 → Reshape 分离时间步 → Concat 交织 → Reshape 展平 → 原始模型 + +使用方式: + python wrap_onnx_multi_input.py + +输入: policy_name.onnx (单输入, input [B, 738] → output [B, 29]) +输出: policy_name_multi.onnx (多输入, 7 个独立 observation term) +""" + +import onnx +from onnx import helper, TensorProto +import numpy as np + +# ============================================================================= +# 1. 定义 Observation Term 拆分方案 (与 physhsi_deploy.yaml 一致) +# ============================================================================= +# use_gym_history: true, history_length: 6 +# C++ ObservationManager 为每个 group 产生 per-term 时间堆叠输出: +# 例如 base_ang_vel: [ang_t5(3)|ang_t4(3)|ang_t3(3)|ang_t2(3)|ang_t1(3)|ang_t0(3)] = 18 dims +# +# 原始模型期望 per-time-step 交织: +# [ang_t5|grav_t5|jpos_t5|jvel_t5|ee_t5|act_t5|task_t5| (t-5: 123) +# ang_t4|grav_t4|jpos_t4|jvel_t4|ee_t4|act_t4|task_t4| (t-4: 123) +# ... ... +# ang_t0|grav_t0|jpos_t0|jvel_t0|ee_t0|act_t0|task_t0] (t: 123) +# +# 顺序必须与用户指定的 738 维组成严格一致: +# 偏移 0-122: t-5 (最新 = h=0 in C++ history) +# 偏移 123-245: t-4 +# 偏移 246-368: t-3 +# 偏移 369-491: t-2 +# 偏移 492-614: t-1 +# 偏移 615-737: t (最旧 = h=5 in C++ history) + +HISTORY_LENGTH = 6 +OBS_TERMS = [ + # (ONNX 输入名称, 单帧维度) + # ★ 顺序必须与用户在每步 123 维中的排列一致 ★ + ("base_ang_vel", 3), # 偏移 0-2 + ("projected_gravity", 3), # 偏移 3-5 + ("joint_pos_rel", 29), # 偏移 6-34 + ("joint_vel_rel", 29), # 偏移 35-63 + ("end_effector_pos", 15), # 偏移 64-78 + ("last_action", 29), # 偏移 79-107 + ("task_obs", 15), # 偏移 108-122 +] +PER_STEP_DIM = sum(dim for _, dim in OBS_TERMS) # 123 +TOTAL_DIM = PER_STEP_DIM * HISTORY_LENGTH # 738 + +assert PER_STEP_DIM == 123, f"每步维度应为 123，实际为 {PER_STEP_DIM}" +assert TOTAL_DIM == 738, f"总维度应为 738，实际为 {TOTAL_DIM}" +print(f"每步维度: {PER_STEP_DIM} ✓") +print(f"总观测维度: {TOTAL_DIM} ✓") + +# ============================================================================= +# 2. 加载原始 ONNX 模型 +# ============================================================================= +INPUT_MODEL = "policy_name.onnx" +OUTPUT_MODEL = "policy_name_multi.onnx" + +original_model = onnx.load(INPUT_MODEL) +print(f"\n加载原始模型: {INPUT_MODEL}") +print(f" 输入: {[inp.name for inp in original_model.graph.input]}") +print(f" 输出: {[out.name for out in original_model.graph.output]}") + +original_graph = original_model.graph +orig_input = original_graph.input[0] # 'input': [B, 738] +orig_output = original_graph.output[0] # 'output': [B, 29] + +# ============================================================================= +# 3. 创建多输入包装图 — 核心: Reshape → Concat → Reshape 实现交织 +# ============================================================================= + +# --- 3.1 创建 7 个新输入 --- +new_inputs = [] +reshape_outputs = [] # 每个输入 Reshape 后的中间结果名称 + +for term_name, dim in OBS_TERMS: + input_dim = dim * HISTORY_LENGTH + new_input = helper.make_tensor_value_info( + term_name, + TensorProto.FLOAT, + ["batch_size", input_dim] + ) + new_inputs.append(new_input) + +# --- 3.2 为每个输入创建 Reshape [B, dim×6] → [B, 6, dim] --- +# 使用 initializer 提供 reshape 的目标形状 +new_nodes = [] +new_initializers = list(original_graph.initializer) # 保留原始权重 + +for i, (term_name, dim) in enumerate(OBS_TERMS): + # Reshape 目标形状: [batch_size, 6, dim] + # ONNX Reshape 支持用 0 表示"保持原维度", -1 表示"推断" + # 这里我们用显式值: shape = [0, 6, dim] 或更安全地使用常量 + shape_name = f"reshape_shape_{i}" + + # 创建形状常量 initializer [3]: [batch_size, 6, dim] + # 用 0 表示 batch_size 维度保持动态 + shape_data = np.array([0, HISTORY_LENGTH, dim], dtype=np.int64) + shape_init = helper.make_tensor( + name=shape_name, + data_type=TensorProto.INT64, + dims=[3], + vals=shape_data.tobytes(), + raw=True + ) + new_initializers.append(shape_init) + + reshape_output = f"reshape_{term_name}" + reshape_node = helper.make_node( + 'Reshape', + inputs=[term_name, shape_name], + outputs=[reshape_output], + name=f'reshape_{term_name}' + ) + new_nodes.append(reshape_node) + reshape_outputs.append(reshape_output) + +# --- 3.3 Concat 所有 Reshape 输出在 axis=2 上 → [B, 6, 123] --- +concat_node = helper.make_node( + 'Concat', + inputs=reshape_outputs, + outputs=['interleaved_obs'], + axis=2, # 在最后一维拼接: 3+3+29+29+15+29+15 = 123 + name='interleave_concat' +) +new_nodes.append(concat_node) + +# --- 3.4 Reshape [B, 6, 123] → [B, 738] --- +# 行主序展平: row0(t-5) | row1(t-4) | ... | row5(t) +final_shape_name = "final_reshape_shape" +final_shape_data = np.array([0, TOTAL_DIM], dtype=np.int64) # [batch_size, 738] +final_shape_init = helper.make_tensor( + name=final_shape_name, + data_type=TensorProto.INT64, + dims=[2], + vals=final_shape_data.tobytes(), + raw=True +) +new_initializers.append(final_shape_init) + +final_reshape = helper.make_node( + 'Reshape', + inputs=['interleaved_obs', final_shape_name], + outputs=['obs_flat_738'], + name='final_reshape' +) +new_nodes.append(final_reshape) + +# ============================================================================= +# 4. 将原始模型的节点追加进来，输入名替换为包装后的输出 +# ============================================================================= +for node in original_graph.node: + new_node = onnx.NodeProto() + new_node.CopyFrom(node) + # 将引用 'input' 的地方替换为 'obs_flat_738' + new_input_names = [] + for inp_name in node.input: + if inp_name == orig_input.name: + new_input_names.append('obs_flat_738') + else: + new_input_names.append(inp_name) + del new_node.input[:] + new_node.input.extend(new_input_names) + new_nodes.append(new_node) + +# ============================================================================= +# 5. 创建输出 +# ============================================================================= +new_output = helper.make_tensor_value_info( + orig_output.name, + TensorProto.FLOAT, + ["batch_size", 29] +) + +# ============================================================================= +# 6. 构建新的 Graph 和 Model +# ============================================================================= +new_graph = helper.make_graph( + nodes=new_nodes, + name="multi_input_interleaved", + inputs=new_inputs, + outputs=[new_output], + initializer=new_initializers, +) + +new_model = helper.make_model( + new_graph, + producer_name="unitree_rl_lab_interleaved_wrapper", + opset_imports=[onnx.helper.make_opsetid("", 11)], + ir_version=6, +) + +# ============================================================================= +# 7. 验证并保存 +# ============================================================================= +onnx.checker.check_model(new_model) +print("\nONNX 模型结构验证通过 ✓") +onnx.save(new_model, OUTPUT_MODEL) +print(f"已保存多输入模型: {OUTPUT_MODEL}") + +# ============================================================================= +# 8. 打印最终结构 +# ============================================================================= +print("\n" + "=" * 60) +print("包装后 ONNX 模型结构") +print("=" * 60) +print(f"\n输入 ({len(new_model.graph.input)} 个):") +for inp in new_model.graph.input: + shape = [d.dim_param if d.dim_param else d.dim_value + for d in inp.type.tensor_type.shape.dim] + print(f" '{inp.name}': shape={shape}") + +print(f"\n输出 ({len(new_model.graph.output)} 个):") +for out in new_model.graph.output: + shape = [d.dim_param if d.dim_param else d.dim_value + for d in out.type.tensor_type.shape.dim] + print(f" '{out.name}': shape={shape}") + +print(f"\n包装节点: {len(new_nodes)} (7 Reshape + 1 Concat + 1 Reshape + 7 原始)") +print(f"初始化器/权重: {len(new_model.graph.initializer)} 个") + +# ============================================================================= +# 9. 运行时验证 — 正确交织 vs 原始单输入 +# ============================================================================= +print("\n" + "=" * 60) +print("运行时推理验证 (正确交织 vs 原始模型)") +print("=" * 60) + +try: + import onnxruntime as ort + + wrapped_session = ort.InferenceSession(OUTPUT_MODEL) + orig_session = ort.InferenceSession(INPUT_MODEL) + + # --- 9.1 生成 per-term 时间堆叠的测试输入 --- + # C++ ObservationManager 为每个 group 产生: + # [t-5_data | t-4_data | t-3_data | t-2_data | t-1_data | t_data] + # 其中 h=0 是最老帧 (t-5), h=5 是最新帧 (t) + per_term_inputs = {} + raw_frames = [] # 按时间步存储各 term 的原始值 + + np.random.seed(42) + for step in range(HISTORY_LENGTH): + frame = {} + for term_name, dim in OBS_TERMS: + frame[term_name] = np.random.randn(dim).astype(np.float32) + raw_frames.append(frame) + + # 构造 C++ 侧 per-term 时间堆叠输入 + # 例如 base_ang_vel: [t-5(3) | t-4(3) | t-3(3) | t-2(3) | t-1(3) | t(3)] = 18 + for term_name, dim in OBS_TERMS: + stacked = [] + for step in range(HISTORY_LENGTH): # h=0 (t-5) to h=5 (t) + stacked.append(raw_frames[step][term_name]) + per_term_inputs[term_name] = np.concatenate(stacked).reshape(1, -1).astype(np.float32) + + # --- 9.2 运行包装后的模型 --- + wrapped_output = wrapped_session.run(None, per_term_inputs) + + # --- 9.3 构造正确的 per-time-step 交织输入与原始模型对比 --- + # 正确顺序: [t-5_all123 | t-4_all123 | t-3_all123 | t-2_all123 | t-1_all123 | t_all123] + correct_flat = [] + for step in range(HISTORY_LENGTH): # t-5 first, t last + for term_name, dim in OBS_TERMS: + correct_flat.append(raw_frames[step][term_name]) + correct_input = np.concatenate(correct_flat).reshape(1, -1).astype(np.float32) + + orig_output = orig_session.run(None, {"input": correct_input}) + + # --- 9.4 对比 --- + diff = np.abs(wrapped_output[0] - orig_output[0]) + print(f"\n 包装模型输出形状: {wrapped_output[0].shape}") + print(f" 原始模型输出形状: {orig_output[0].shape}") + print(f" 最大误差: {diff.max():.6e}") + print(f" 均方误差: {np.mean(diff**2):.6e}") + + if diff.max() < 1e-5: + print(" ✓ 输出完全一致 — 交织逻辑正确") + else: + print(f" ✗ 存在差异!") + + # --- 9.5 额外验证: 错误排序对比 (确认旧版 bug) --- + print("\n" + "-" * 40) + print("对比: 如果用逐项分组 (旧版错误做法) 会怎样?") + wrong_flat = [] + for term_name, dim in OBS_TERMS: + for step in range(HISTORY_LENGTH): + wrong_flat.append(raw_frames[step][term_name]) + wrong_input = np.concatenate(wrong_flat).reshape(1, -1).astype(np.float32) + wrong_output = orig_session.run(None, {"input": wrong_input}) + wrong_diff = np.abs(wrong_output[0] - orig_output[0]) + print(f" 错误排序 → 原始模型输出的最大误差: {wrong_diff.max():.6e}") + print(f" 错误排序 → 原始模型输出的均方误差: {np.mean(wrong_diff**2):.6e}") + if wrong_diff.max() > 0.01: + print(" ✓ 确认旧版 wrapper 会产生完全错误的输出!") + +except ImportError: + print("\n (跳过运行时验证: onnxruntime 未安装)") +except Exception as e: + import traceback + print(f"\n ⚠ 运行时验证失败: {e}") + traceback.print_exc()