From 9d4a8ae6b0b40de6c249420e8475dc95bcaf2534 Mon Sep 17 00:00:00 2001 From: m0_51228868 Date: Mon, 8 Jun 2026 18:19:14 +0800 Subject: [PATCH 1/6] align wmp data semantics --- README.md | 180 +++-- WMP_NETWORK_DESIGN_DIFF_REPORT.txt | 678 ++++++++++++++++++ legged_lab/envs/a1/a1_config.py | 67 +- legged_lab/envs/base/base_config.py | 24 +- legged_lab/envs/base/base_env.py | 449 +++++++++++- .../convert_original_wmp_checkpoint.py | 129 ++++ .../sensors/wmp_partial_tiled_camera.py | 146 +++- legged_lab/utils/env_utils/scene.py | 11 +- legged_lab/world_models/wmp/controller.py | 202 ++++-- 9 files changed, 1685 insertions(+), 201 deletions(-) create mode 100644 WMP_NETWORK_DESIGN_DIFF_REPORT.txt create mode 100644 legged_lab/scripts/convert_original_wmp_checkpoint.py diff --git a/README.md b/README.md index 9a8a23b..011f670 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,8 @@ - 保持原始 LeggedLab 的直接环境工作流 - 便于按个人需求快速修改训练脚本、日志和任务配置 - 支持使用 `wandb` 作为默认训练日志后端 -- 增加 B2 RGBD / WMP-AMP 训练链路与检查脚本 +- 增加 A1 WMP-AMP 训练链路、原版 WMP 对齐检查与辅助脚本 +- 保留 B2 RGBD / WMP world model 前向检查脚本 - 增加 RB160W 轮腿机器人任务、混合腿位置/轮速度控制与模型验证脚本 ## Fork Statement @@ -93,7 +94,7 @@ python legged_lab/scripts/launcher_gui.py ## Available Tasks 当前注册任务包括： -- `a1_amp_flat` +- `a1_amp_flat`, `a1_wmp_amp_terrain` - `h1_flat`, `h1_rough` - `g1_flat`, `g1_rough` - `gr2_flat`, `gr2_rough` @@ -160,25 +161,106 @@ python legged_lab/scripts/inspect_rb160w_model.py --task=rb160w_flat --mode=swee 本仓库沿用 IsaacLab / rsl_rl 的并行训练方式： - https://isaac-sim.github.io/IsaacLab/main/source/features/multi_gpu.html -## WMP RSSM World Model Smoke Test +## A1 WMP-AMP -当前分支已开始移植 WMP（World Model-based Perception for Visual Legged Locomotion）的 RSSM 世界模型到 -`b2_rgbd`。本阶段只提供模型前向检查，不接入 PPO 联训，也不改变现有 `b2_rgbd_*` 训练观测。 +当前主线 WMP 任务是 `a1_wmp_amp_terrain`，目标是在 IsaacLab + rsl_rl 5 中复现 ByteDance WMP 的 A1 地形训练链路。它不是对原仓库的逐文件搬运，而是把原版 WMP 的 actor/critic/RSSM/DepthPredictor/AMP 语义接入当前 LeggedLab 结构。 -当前链路： +核心文件： +- 任务配置：`legged_lab/envs/a1/a1_config.py` +- runner：`legged_lab/runners/wmp_amp_runner.py` +- PPO/AMP 算法：`legged_lab/algorithms/wmp_amp_ppo.py` +- WMP 模型控制器：`legged_lab/world_models/wmp/controller.py` +- WMP policy/critic 模型：`legged_lab/models/wmp_mlp_model.py` +- partial depth camera：`legged_lab/sensors/wmp_partial_tiled_camera.py` + +训练命令： + +```bash +python legged_lab/scripts/train.py \ + --task=a1_wmp_amp_terrain \ + --num_envs=4096 \ + --headless \ + --enable_cameras \ + --runner=wmp_amp \ + --logger=wandb +``` + +小规模 smoke： + +```bash +python legged_lab/scripts/train.py \ + --task=a1_wmp_amp_terrain \ + --num_envs=2 \ + --headless \ + --enable_cameras \ + --runner=wmp_amp \ + --max_iterations=1 \ + --num_steps_per_env=2 \ + --num_mini_batches=1 \ + --logger=tensorboard +``` + +不开相机的 Python 链路 dry smoke： + +```bash +python legged_lab/scripts/train.py \ + --task=a1_wmp_amp_terrain \ + --num_envs=2 \ + --headless \ + --runner=wmp_amp \ + --max_iterations=1 \ + --num_steps_per_env=2 \ + --num_mini_batches=1 \ + --logger=tensorboard +``` + +说明：`--runner=wmp_amp` 且不传 `--enable_cameras` 时，`train.py` 会关闭 RGBD camera 并启用 zero-depth fallback；这只适合检查 Python 训练链路，不适合评估 WMP 视觉能力。 + +### 与原版 WMP 的当前对齐状态 + +- Actor 主干与原版对齐：`history_encoder(210 -> 35) + command(3) + wmp_encoder(512 -> 32) -> MLP(70 -> 12)`。 +- Critic 输入维度已按原版恢复：`critic_obs(285) + wmp_latent(32) = 317`，第一层对应原版 `critic.0.weight: (512, 317)`。 +- WMP/RSSM 默认使用 `deter=512` 作为 policy feature；`update_interval=5`，因此 WMP action 维度为 `12 * 5 = 60`。 +- WMP proprioception 为 33 维：`ang_vel(3) + projected_gravity(3) + command(3) + joint_pos(12) + joint_vel(12)`。 +- DepthPredictor 使用 forward height map 525 维和 prop 33 维，输出 `64x64x1` WMP depth image。 +- forward height map 已按原版修正为前方 `x=[0,2], y=[-1.2,1.2]` 采样语义；IsaacLab centered grid 通过 `forward_offset=(1,0,0)` 得到 `x'=x+1`。 +- AMP discriminator 输入为 60 维，即 `AMP obs(30) + next AMP obs(30)`；motion 默认来自 `datasets/wmp_mocap_motions/*.txt`。 + +原版 critic layout 在当前环境中构造为： ```text -b2_rgbd Gemini2 depth - -> clean/clip/normalize - -> resize to 1x64x64 (NCHW) - -> convert to WMP image format 64x64x1 (NHWC) - -> WMP MultiEncoder - -> RSSM obs_step - -> deter feature(512) / full feature(1536) - -> MultiDecoder image reconstruction smoke test +contact_flag(8) ++ contact_force(12) ++ d_gain_scale(12) ++ p_gain_scale(12) ++ com_pos(3) ++ added_mass(1) ++ restitution(1) ++ friction(1) ++ base_lin_vel(3) ++ actor_obs(45) += 98 + +98 + height_scan(187) = critic_obs(285) +critic_obs(285) + wmp_latent(32) = critic_input(317) ``` -运行检查： +### 仍需注意的差异 + +当前代码已经修正了旧报告中提到的 `critic 271` 问题；如果看到 `WMP_NETWORK_DESIGN_DIFF_REPORT.txt` 中相关描述，应以当前源码和本 README 为准。 +A1 WMP privileged scale 也已恢复到原版：`contact_force=0.005`、`pd_gains=5.0`、`com_pos=20.0`。 +当前从零训练不要求兼容旧 `.pt`，actor obs、critic obs、WMP prop、AMP obs 和 action 使用 IsaacLab runtime joint order 保持内部自洽。 + +仍可能影响训练效果的差异： +- IsaacLab 与原版 IsaacGym/legged_gym 的仿真、接触、相机和资产加载不同。 +- 原版 checkpoint 的 joint order 与当前 runtime joint order 不直接兼容，需要额外转换。 +- WMP replay 由 `WMPFixedEpisodeReplayBuffer` 管理，不是原版 `wm_dataset/wm_buffer` 的逐行实现。 +- partial camera 和 depth buffer 使用 IsaacLab sensor/render 调度，depth 帧时序需要重点检查。 +- AMP motion 经过 `A1CanonicalRetargetAdapter` 映射到当前 env joint order，专家数据均值和当前机器人初态应在启动日志中核对。 + +### WMP 检查脚本 + +World model 前向检查： ```bash python legged_lab/scripts/inspect_wmp_world_model.py \ @@ -196,6 +278,24 @@ python legged_lab/scripts/inspect_wmp_world_model.py \ - `RSSM full feature shape=(1, 1536)` - `decoded image shape=(1, 64, 64, 1)` +地形与边缘 mask 检查： + +```bash +python legged_lab/scripts/inspect_wmp_terrain_generator.py --headless +python legged_lab/scripts/inspect_wmp_edge_mask.py --task=a1_wmp_amp_terrain --headless --num_envs=4 +``` + +原版 checkpoint 转换： + +```bash +python legged_lab/scripts/convert_original_wmp_checkpoint.py \ + /path/to/original/model.pt \ + --reference /path/to/leggedlab/reference.pt \ + --output /path/to/converted.pt +``` + +转换脚本会迁移 actor、world model、DepthPredictor；critic、optimizer、AMP 状态保留 reference 模板，主要用于推理/播放检查。 + ## Troubleshooting ### GitHub 大文件 warning @@ -298,53 +398,3 @@ train.py -> get_cfgs -> BaseEnv() PPO update log/save ``` - -## WMP + AMP-PPO - -当前已加入 `b2_rgbd_wmp_amp_flat` 和 `b2_rgbd_wmp_amp_terrain`：在 rsl_rl 5 PPO 基础上拼接 WMP RSSM feature，并用 WMP AMP 判别器替换 rollout reward。该版本用于打通训练链路，不承诺复现论文最终性能。 - -- WMP 输入：Gemini2 `distance_to_image_plane -> clamp/normalize -> 1x64x64 -> NHWC image`，policy feature 默认使用 RSSM `deter=512`。 -- AMP 数据：默认使用 `datasets/wmp_mocap_motions/*.txt`，来源于 ByteDance WMP 的四足 mocap；B2 joint/foot 顺序会在启动时打印，必要时需要进一步做显式 remap。 -- RGB：第一版不使用 RGB，相机任务只启用 depth。 -- DepthPredictor：第一版暂不接入。 - -快速 dry smoke（不启真实相机，使用零 depth fallback，只用于检查 PPO/AMP/WMP Python 链路）： - -```bash -python legged_lab/scripts/train.py \ - --task=b2_rgbd_wmp_amp_flat \ - --num_envs=2 \ - --headless \ - --runner=wmp_amp \ - --max_iterations=1 \ - --num_steps_per_env=2 \ - --num_mini_batches=1 \ - --logger=tensorboard -``` - -真实 depth 训练需要开启 camera： - -```bash -python legged_lab/scripts/train.py \ - --task=b2_rgbd_wmp_amp_flat \ - --num_envs=64 \ - --headless \ - --enable_cameras \ - --runner=wmp_amp \ - --max_iterations=10000 \ - --logger=wandb \ - --log_project_name=b2_rgbd_wmp_amp \ - --run_name=wmp_amp_depth_only_v1 -``` - -地形版 WMP-AMP： - -```bash -python legged_lab/scripts/train.py \ - --task=b2_rgbd_wmp_amp_terrain \ - --num_envs=64 \ - --headless \ - --enable_cameras \ - --runner=wmp_amp \ - --logger=wandb -``` diff --git a/WMP_NETWORK_DESIGN_DIFF_REPORT.txt b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt new file mode 100644 index 0000000..5b4415a --- /dev/null +++ b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt @@ -0,0 +1,678 @@ +WMP 网络设计对比报告 +原始项目路径: /tmp/bytedance_WMP +当前项目路径: /home/tower/Bags/LeggedLab +更新时间 : 2026-06-08 +关注问题 : 当前项目和原版 WMP 之间，训练网络、训练接口、数据语义和输入排列的主要差异 + + +生成口径 +-------- +本报告以当前工作区源码为准，并结合 isaaclab conda 环境下的运行时检查结果。 +检查命令使用: + + source /home/tower/miniconda/etc/profile.d/conda.sh + conda activate isaaclab + /home/tower/Bags/IsaacLab/isaaclab.sh -p /tmp/inspect_leggedlab_wmp_runtime.py --headless --num_envs 2 + +注意: + 旧版报告中提到“当前 critic 为 271 维，不等价原版 317 维”的结论已经过期。 + 当前 A1 WMP critic 已恢复为: + + critic_obs(285) + wmp_latent(32) = critic_input(317) + + +1. 总结结论 +=========== + +1. Actor 主干网络在 shape 上对齐原版。 + 原版和当前都是: + history_encoder: 210 -> 256 -> 128 -> 35 + WMP encoder : 512 -> 64 -> 64 -> 32 + actor MLP : 70 -> 256 -> 128 -> 64 -> 12 + 其中: + 70 = history_latent(35) + command(3) + wmp_latent(32) + +2. Critic 当前已经恢复到原版 WMP 的输入维度设计。 + 原版: + critic_obs = privileged_obs = 285 + critic_input = 285 + wmp_latent(32) = 317 + critic.0.weight = (512,317) + 当前: + BaseEnv 构造 original WMP critic layout，critic_obs = 285 + WMPMLPModel critic 不走 history encoder，拼接 wmp_latent 后输入 317 + mlp.0.weight = (512,317) + +3. 本轮已修复两个关键数据语义问题。 + - forward height scanner 已前移 1m。 + IsaacLab GridPattern 以局部原点为中心生成 x=[-1,1], y=[-1.2,1.2]； + 加 offset=(1,0,0) 后得到 x'=x+1=[0,2]，与原版 WMP forward height map 对齐。 + - A1 WMP privileged obs scale 已恢复到原版: + contact_force = 0.005 + pd_gains = 5.0 + com_pos = 20.0 + +4. 当前从零训练的输入排列在项目内部自洽。 + actor obs、critic obs、WMP prop、AMP obs、action 都使用 IsaacLab runtime joint order。 + AMP expert motion 通过 A1CanonicalRetargetAdapter(target_joint_order="env") 映射到 env joint order。 + 因此本轮目标不是兼容旧 pt，而是保证当前训练链路内部一致。 + +5. World Model/RSSM 架构参数基本对齐原版。 + 当前 A1 WMP 配置: + update_interval=5 + train_start_steps=10000 + train_steps_per_iter=10 + batch_size=16 + batch_length=64 + model_lr=1e-4 + feature_type="deter" + 但 replay、optimizer、训练时序、real depth 覆盖路径仍不是原版逐行实现。 + +6. DepthPredictor 和 AMP discriminator 的网络 shape 对齐。 + DepthPredictor 输入 forward_height_map(525) + prop(33) = 558，输出 64x64x1。 + AMP discriminator 输入为 AMP obs(30) + next AMP obs(30) = 60。 + +7. 目前最关键的剩余风险不再是 critic 维度、forward map 范围或 privileged scale，而是: + - IsaacLab / IsaacGym 的物理、接触、相机和资产语义差异； + - WMP replay 与原版 wm_dataset/wm_buffer 的采样语义差异； + - depth camera 帧时序和 partial camera 映射； + - AMP retarget 后 expert 数据分布； + - rsl_rl v5 PPO/storage/distribution/normalization 接口差异。 + + +2. 运行时数据语义检查 +===================== + +2.1 Joint order +--------------- +当前 A1 USD 在 IsaacLab runtime 中的 joint order: + + ['FL_hip_joint', 'FR_hip_joint', 'RL_hip_joint', 'RR_hip_joint', + 'FL_thigh_joint', 'FR_thigh_joint', 'RL_thigh_joint', 'RR_thigh_joint', + 'FL_calf_joint', 'FR_calf_joint', 'RL_calf_joint', 'RR_calf_joint'] + +原版 URDF / AMP motion 源顺序: + + FR_hip, FR_thigh, FR_calf, + FL_hip, FL_thigh, FL_calf, + RR_hip, RR_thigh, RR_calf, + RL_hip, RL_thigh, RL_calf + +结论: + 当前不要求旧 pt 兼容，因此允许 runtime joint order 与原版源顺序不同。 + 当前训练链路内部使用同一个 env joint order，AMP expert motion 在加载时 retarget 到 env order。 + 这对从零训练是自洽的；若加载原版 policy/world model/AMP checkpoint，则仍需要权重或输入顺序转换。 + + +2.2 输入排列 +------------ +当前 runtime shape: + + policy : (2, 412) + critic : (2, 285) + current_actor : (2, 45) + current_critic : (2, 98) + wmp_prop : (2, 33) + amp_obs : (2, 30) + forward_height : (2, 525) + +actor current obs layout: + base_ang_vel(3) + + projected_gravity(3) + + command(3) + + joint_pos(12) + + joint_vel(12) + + last_action(12) + = 45 + +WMP prop layout: + base_ang_vel(3) + + projected_gravity(3) + + command(3) + + joint_pos(12) + + joint_vel(12) + = 33 + +AMP obs layout: + joint_pos(12) + + base_lin_vel_b(3) + + base_ang_vel_b(3) + + joint_vel(12) + = 30 + +结论: + joint_pos、joint_vel、action 的 12 维排列均来自 env.robot.joint_names。 + AMP expert motion 已重排到同一 env order，因此 actor / WMP / AMP / action 在当前训练中内部一致。 + + +2.3 Contact body order +---------------------- +运行时 contact sensor body_names: + + ['trunk', 'FL_hip', 'FL_thigh', 'FL_calf', 'FL_foot', + 'FR_hip', 'FR_thigh', 'FR_calf', 'FR_foot', + 'RL_hip', 'RL_thigh', 'RL_calf', 'RL_foot', + 'RR_hip', 'RR_thigh', 'RR_calf', 'RR_foot'] + +feet contact body_names: + + ['FL_foot', 'FR_foot', 'RL_foot', 'RR_foot'] + +wmp privileged contact body_names: + + ['FL_thigh', 'FL_calf', 'FR_thigh', 'FR_calf', + 'RL_thigh', 'RL_calf', 'RR_thigh', 'RR_calf'] + +结论: + 当前 contact_flag(8) 来自 thigh/calf，contact_force(12) 来自 feet forces 并 reshape。 + body 顺序与 IsaacLab runtime 一致；与原版 IsaacGym 接触力数值分布仍可能不同。 + + +2.4 DelayBuffer / last action +----------------------------- +运行时检查: + + time_lags: [0, 0] + buffer_return == input_action + actor obs action slice 为当前输入 action + +结论: + delay lag=0 时，IsaacLab DelayBuffer.compute(data) 先写入当前 data，再返回当前 data。 + 当前 last_action 语义与无延迟训练链路自洽。 + + +2.5 Forward height map +---------------------- +原版 A1 WMP: + + measured_forward_points_x = [0.0, ..., 2.0] + measured_forward_points_y = [-1.2, ..., 1.2] + +当前修复后 runtime 前几个采样点: + + [[0.0, -1.2], [0.1, -1.2], ..., [1.1, -1.2]] + +结论: + forward_height_scanner 已通过 forward_offset=(1,0,0) 对齐到原版 x=[0,2], y=[-1.2,1.2] 语义。 + 普通 height_scan 仍保持原本 centered grid，用于 critic/policy 地形高度。 + + +3. Actor 网络对比 +================= + +3.1 原版 ActorCriticWMP +----------------------- +文件: + /tmp/bytedance_WMP/rsl_rl/modules/actor_critic_wmp.py + +结构: + history_encoder: + Linear(210,256) -> ELU + Linear(256,128) -> ELU + Linear(128,35) + + wm_feature_encoder: + Linear(512,64) -> ELU + Linear(64,64) -> ELU + Linear(64,32) + + actor: + Linear(70,256) -> ELU + Linear(256,128) -> ELU + Linear(128,64) -> ELU + Linear(64,12) + +输入: + latent_vector = history_encoder(history) + command = observations[:, privileged_dim + 6 : privileged_dim + 9] + wm_latent_vector = wm_feature_encoder(wm_feature) + actor_input = concat(latent_vector, command, wm_latent_vector) + + +3.2 当前 WMPMLPModel actor +-------------------------- +文件: + legged_lab/models/wmp_mlp_model.py + legged_lab/envs/a1/a1_config.py + +结构: + 与原版 actor 主干 shape 对齐: + history_encoder: 210 -> 256 -> 128 -> 35 + wmp_encoder: 512 -> 64 -> 64 -> 32 + actor MLP: 70 -> 256 -> 128 -> 64 -> 12 + +当前输入构造: + base_obs = concat(non_wmp obs groups) + history_obs, extra = split_history_and_extra(base_obs) + history = remove_command_from_each_step(history_obs) + command = command_from_last_history_step(history_obs) + wmp_latent = wmp_encoder(obs["wmp"]) + actor_input = concat(history_encoder(history), command, wmp_latent) + +差异: + 1. 原版 runner 显式传入 obs/history/wm_feature。 + 当前通过 rsl_rl v5 TensorDict 传入 policy/wmp groups。 + 2. 原版 history 由 WMPRunner.trajectory_history 维护。 + 当前 history 来自 BaseEnv.actor_obs_buffer，再由 WMPMLPModel reshape。 + 3. 原版 actor 没有 latent obs_normalizer。 + 当前是否启用 normalizer 取决于 rsl_rl v5 cfg 映射后的 obs_normalization。 + 4. joint order 对当前从零训练自洽；对原版 checkpoint 不直接兼容。 + + +4. Critic 网络对比 +================== + +4.1 原版 critic +--------------- +原版 ActorCriticWMP critic: + critic_wm_feature_encoder: + Linear(512,64) -> ELU + Linear(64,64) -> ELU + Linear(64,32) + + critic MLP: + Linear(317,512) -> ELU + Linear(512,256) -> ELU + Linear(256,128) -> ELU + Linear(128,1) + +实际原版第一层: + critic.0.weight: (512,317) + +含义: + 317 = 原版 privileged critic obs 285 + WMP latent 32。 + + +4.2 当前 critic +--------------- +当前 A1 WMP 环境启用: + self.robot.wmp_privileged_contact_body_names = [".*thigh.*", ".*calf.*"] + +因此 BaseEnv 走 original WMP critic obs 路径: + legged_lab/envs/base/base_env.py + _compute_original_wmp_critic_obs() + +当前 critic obs layout: + contact_flag(8) + + contact_force(12) + + d_gain_scale(12) + + p_gain_scale(12) + + com_pos(3) + + added_mass(1) + + restitution(1) + + friction(1) + + base_lin_vel(3) + + actor_obs(45) + = 98 + + 98 + height_scan(187) = 285 + +WMPMLPModel critic: + use_history_encoder=True 仅让 actor 使用 history path。 + critic_uses_history_encoder 默认 False，因此 critic 直接使用: + latent = concat(critic_obs(285), wmp_encoder(wmp_feature)(32)) + latent_dim = 317 + +当前 critic MLP: + Linear(317,512) -> ELU + Linear(512,256) -> ELU + Linear(256,128) -> ELU + Linear(128,1) + + +4.3 Critic 结论 +--------------- +当前 critic 的网络输入维度已经和原版对齐。 + +已解决的旧问题: + - 不再是 271 维 critic input。 + - 不再需要把“恢复 critic 到 317 维”列为首要待办。 + - privileged scale 已恢复原版 contact_force=0.005、pd_gains=5.0、com_pos=20.0。 + +仍需关注: + 1. contact_force/contact_flag 来自 IsaacLab contact sensor，可能和原版 IsaacGym 接触力分布不同。 + 2. friction/restitution/mass/COM/gain 来自当前记录型 domain randomization buffer，并会尝试从 sim 刷新。 + 3. actor_obs(45) 中 joint_pos/joint_vel/last_action 的 joint order 对当前训练自洽，但不直接兼容原版 pt。 + 4. height_scan 的数值分布仍受 IsaacLab RayCaster 与 terrain 实现影响。 + + +5. World Model / RSSM 网络对比 +============================== + +5.1 原版 Dreamer/WMP 配置 +------------------------ +核心配置: + dyn_hidden=512 + dyn_deter=512 + dyn_stoch=32 + dyn_discrete=32 + dyn_rec_depth=1 + dyn_mean_act=none + dyn_std_act=sigmoid2 + dyn_min_std=0.1 + unimix_ratio=0.01 + initial=learned + units=512 + act=SiLU + norm=True + grad_heads=["decoder","reward"] + dyn_scale=0.5 + rep_scale=0.1 + kl_free=1.0 + model_lr=1e-4 + batch_size=16 + batch_length=64 + train_start_steps=10000 + train_steps_per_iter=10 + +RSSM: + prop_dim=33 + action_dim=12*5=60 + deter feature=512 + full feature=1536 + 原版 policy 默认使用 deter feature 512。 + + +5.2 当前 WorldModel 配置 +----------------------- +文件: + legged_lab/world_models/wmp/config.py + legged_lab/world_models/wmp/models.py + legged_lab/world_models/wmp/networks.py + legged_lab/world_models/wmp/controller.py + +当前 A1 WMP 覆盖: + feature_type="deter" + update_interval=5 + train_start_steps=10000 + train_steps_per_iter=10 + batch_size=16 + batch_length=64 + model_lr=1e-4 + replay_capacity_episodes=50000 + replay_device="cpu" + +shape 对齐点: + image encoder 输入: 64x64x1 + prop encoder 输入: 33 + action_dim: 60 + deter feature: 512 + full feature: 1536 + decoder image 输出: 64x64x1 + decoder prop 输出: 33 + + +5.3 World Model 差异结论 +------------------------ +网络 shape 和主要超参已经接近原版。 + +当前本轮同步的训练/深度语义: + 1. train_start_steps 现在按 replay.num_steps 判断，而不是 total_env_steps。 + 这更接近原版基于 world model 数据量启动训练的语义。 + 2. partial real depth 优先通过 get_wmp_depth_observations() 获取。 + BaseEnv 维护 2 帧 WMP depth buffer，并按原版 depth_buffer[:, -2] 语义返回上一帧。 + 3. WMPPartialTiledCamera 要求实际 image slot order 与 depth_index 完全一致。 + 若 IsaacLab prim order 改变，会直接报错，避免静默错位。 + 4. depth_index_without_crawl_tilt 从非 tilt/crawl 前方地形池采样，并把 tilt/crawl env 拼在后面。 + +训练接口仍不同: + 1. 原版使用 dreamer.tools.Optimizer，当前使用 torch.optim.Adam。 + 2. 原版 runner 内维护 wm_dataset/wm_buffer。 + 当前使用 WMPFixedEpisodeReplayBuffer。 + 3. 原版在 WMPRunner 中直接覆盖: + wm_obs["image"][env.depth_index] = infos["depth"] + 当前通过 WMPTrainingController + BaseEnv depth buffer/partial camera 获取真实 depth。 + 4. 原版按 env.global_counter % update_interval 驱动 RSSM tick。 + 当前由 WMPTrainingController.step_counter 和 pending_obs 管理。 + 5. is_first、episode finish、action_history flatten、reward accumulation 的时序仍建议持续做统计检查。 + + +6. DepthPredictor 网络对比 +========================== + +6.1 原版 DepthPredictor +----------------------- +输入: + forward_heightmap_dim=525 + prop_dim=33 + concat input=558 + +结构: + MLP: + Linear(558,256) -> ELU + Linear(256,128) -> ELU + Linear(128,4096) + + deconv decoder: + ConvTranspose2d(256,128,4,2) + ImgChLayerNorm + ELU + ConvTranspose2d(128,64,4,2) + ImgChLayerNorm + ELU + ConvTranspose2d(64,32,4,2) + ImgChLayerNorm + ELU + ConvTranspose2d(32,1,4,2) + + +6.2 当前 DepthPredictor +----------------------- +文件: + legged_lab/world_models/wmp/depth_predictor.py + +结构: + 与原版 shape 对齐。 + +当前实现差异: + 1. forward 支持 leading batch/time 维。 + 2. 训练数据来自 WMPFixedEpisodeReplayBuffer.sample_real_depth()。 + 3. 当前真实 depth 由 IsaacLab TiledCamera / WMPPartialTiledCamera 提供。 + 4. 当前 depth 预处理处理 IsaacLab 正深度，并转换成 WMP 约定的 64x64x1 图。 + +结论: + DepthPredictor 网络本身不是主要问题。 + 本轮已修正 forward_height_map x 方向语义；剩余关键是真实 depth 的帧时序、相机姿态和样本分布。 + + +7. AMP Discriminator 网络对比 +============================= + +7.1 原版 AMPDiscriminator +------------------------- +结构: + input_dim = 30 * 2 = 60 + trunk: + Linear(60,1024) -> ReLU + Linear(1024,512) -> ReLU + head: + Linear(512,1) + +reward 公式: + r_amp = amp_reward_coef * clamp(1 - 1/4 * (D(s,s') - 1)^2, min=0) + + +7.2 当前 AMPDiscriminator +------------------------- +文件: + legged_lab/amp/discriminator.py + legged_lab/amp/motion_loader.py + legged_lab/amp/retarget.py + +结构: + 与原版 shape 对齐: + trunk.0.weight: (1024,60) + trunk.2.weight: (512,1024) + amp_linear.weight: (1,512) + +当前配置: + canonical_obs_dim=30 + target_joint_order="env" + retarget_adapter=A1CanonicalRetargetAdapter + num_preload_transitions=2000000 + reward_coef=0.01 + task_reward_lerp=0.3 + replay_buffer_size=1000000 + grad_penalty_coef=1.0 + +结论: + 网络 shape 对齐。 + expert AMP obs 的源 joint order 与 env runtime joint order 不同，但当前 loader 通过 retarget_adapter 映射到 env order。 + Runner 启动时应继续核对 expert retarget 后的 joint mean/min/max 与当前 env mean。 + + +8. PPO / Distribution / Normalization 差异 +========================================== + +8.1 原版 PPO/AMPPPO +------------------- +原版: + - ActorCriticWMP 同时持有 actor/critic。 + - std 是 ActorCriticWMP 内部参数。 + - 没有 actor latent obs_normalizer。 + - AMPPPO 直接持有 actor_critic/discriminator/AMP replay。 + - RolloutStorage 显式存储 history 和 wm_feature。 + + +8.2 当前 rsl_rl v5 PPO/WMPAMPPPO +-------------------------------- +当前: + - actor 和 critic 是 rsl_rl v5 model factory 分开构建。 + - actor 使用 GaussianDistribution，std 参数为 distribution.std_param。 + - WMPMLPModel 通过 TensorDict obs groups 读取 policy/critic/wmp。 + - WMPAMPPPO 继承 rsl_rl PPO，再 attach AMP discriminator 和 AMP replay。 + - storage/mini-batch generator/normalization 更新逻辑来自 rsl_rl v5。 + +当前 A1 WMP PPO 配置: + actor_hidden_dims=[256,128,64] + critic_hidden_dims=[512,256,128] + learning_rate=3e-4 + schedule="fixed" + desired_kl=0.006 + entropy_coef=0.01 + num_learning_epochs=5 + num_mini_batches=4 + normalize_advantage_per_mini_batch=True + vel_predict_coef=1.0 + num_steps_per_env=24 + +影响: + 即使网络 shape 对齐，rsl_rl v5 的 model/storage/distribution/update 接口仍会让优化轨迹和原版不同。 + + +9. 当前网络与原版网络的对齐程度评级 +==================================== + +Actor 主干: + 对齐程度: 高 + 主要风险: + 输入构造、normalizer、history 时序、原版 checkpoint joint order 兼容性。 + +Critic: + 对齐程度: 中高 + 原因: + 输入维度已恢复为 317，网络层 shape 与原版同构，privileged scale 已恢复原版。 + 主要风险: + contact/domain randomization 的 IsaacLab 数值语义。 + +World Model / RSSM: + 对齐程度: 中高 + 原因: + 网络 shape 和主要超参对齐。 + 主要风险: + replay、optimizer、训练数据时序、real depth 覆盖、is_first 标记。 + +DepthPredictor: + 对齐程度: 高 + 原因: + shape 和结构基本一致，forward height map 采样范围已修正。 + 主要风险: + 真实 depth 分布、相机姿态和帧时序。 + +AMP Discriminator: + 对齐程度: 高 + 原因: + 网络结构和 reward 公式基本一致。 + 主要风险: + AMP obs retarget 后的专家数据分布。 + +PPO/训练接口: + 对齐程度: 中低 + 原因: + rsl_rl v5 改变了模型包装、storage、distribution、normalization 和 update 接口。 + + +10. 后续优先排查建议 +==================== + +第一优先级: 训练启动时断言 +-------------------------- +建议保留或新增: + - critic_obs shape == 285。 + - critic_input shape == 317。 + - wmp_prop shape == 33。 + - amp_obs shape == 30。 + - forward_height_map shape == 525。 + + +第二优先级: critic 字段统计 +-------------------------- +critic 维度已经对齐，下一步不是改维度，而是验收字段语义: + - contact_flag/contact_force 的 body 顺序和均值范围。 + - friction/restitution/mass/COM/gain buffer 与真实 sim 值同步。 + - height_scan 的均值、范围、ray 顺序与原版接近程度。 + + +第三优先级: WMP replay 和 depth 帧时序 +------------------------------------- +重点检查: + - observe_before_policy() 消费的是哪一个 WMP obs。 + - after_env_step() 写入 replay 的 prop/image/action/reward 是否同一 WMP tick。 + - action_history flatten 顺序是否为 [a[t-k+1], ..., a[t]]。 + - is_first 在 episode reset 和 batch sample 中是否与原版一致。 + - real depth 覆盖的 env ids 是否和 partial camera image slot 一致。 + - depth buffer 返回帧是否与原版 depth_buffer[:, -2] 语义一致。 + + +第四优先级: AMP obs / retarget 对齐 +---------------------------------- +建议记录: + - expert AMP obs mean/min/max。 + - policy AMP obs mean/min/max。 + - retarget 后 joint order。 + - discriminator policy/expert score 分布。 + - amp_reward/task_reward 比例。 + + +第五优先级: rsl_rl v5 训练接口 ablation +-------------------------------------- +建议做: + - actor obs_normalization 开关实验。 + - critic obs_normalization 开关实验。 + - normalize_advantage_per_mini_batch 开关实验。 + - desired_kl/schedule 与原版配置对照实验。 + + +11. 最终判断 +============ + +当前项目和原版 WMP 的网络 shape 已经高度接近: + - actor 已对齐； + - critic 输入维度已恢复为 317； + - RSSM/WorldModel shape 与主要超参已对齐； + - DepthPredictor shape 对齐； + - AMP discriminator shape 对齐。 + +本轮针对数据来源语义和输入排列的检查结论: + - forward height map 已对齐原版 [0,2] x [-1.2,1.2] 采样语义； + - privileged scale 已恢复原版 contact_force/pd_gains/com_pos； + - actor / critic / WMP prop / AMP / action 在当前 env runtime joint order 下内部自洽； + - DelayBuffer lag=0 时 last_action 使用当前输入 action，当前无延迟语义自洽； + - partial real depth 现在要求 image slot order 与 depth_index 一致，避免静默错位。 + +因此，后续如果训练效果仍不对应，不能再把“critic 维度不一致”“forward map 仍是 [-1,1]” +或“privileged scale 未恢复原版”作为首要解释。更可能的原因是: + observation 数值分布； + WMP tick / replay / depth buffer 时序； + IsaacLab 与 IsaacGym 的仿真和传感器差异； + AMP retarget 后 expert 数据分布； + rsl_rl v5 PPO/storage/normalization 更新接口差异。 diff --git a/legged_lab/envs/a1/a1_config.py b/legged_lab/envs/a1/a1_config.py index f727d28..d33e62c 100644 --- a/legged_lab/envs/a1/a1_config.py +++ b/legged_lab/envs/a1/a1_config.py @@ -62,7 +62,7 @@ def __post_init__(self): self.scene.terrain_type = "plane" self.scene.terrain_generator = None self.scene.height_scanner.enable_height_scan = False - self.scene.gemini2_camera.enable = False + self.scene.rgbd_camera.enable = False self.scene.env_spacing = 2.5 self.robot.feet_body_names = [".*foot.*"] self.robot.terminate_contacts_body_names = [".*trunk.*"] @@ -161,7 +161,7 @@ class A1WMPAMPRewardCfg(A1AMPRewardCfg): params={"sensor_cfg": SceneEntityCfg("contact_sensor", body_names="(?!.*foot.*).*"), "threshold": 0.1}, ) feet_stumble = RewTerm( - func=mdp.feet_stumble, + func=mdp.wmp_feet_stumble, weight=-0.1, params={"sensor_cfg": SceneEntityCfg("contact_sensor", body_names=[".*foot.*"])}, ) @@ -201,36 +201,42 @@ def __post_init__(self): self.scene.terrain_type = "generator" self.scene.terrain_generator = WMP_MIXED_TERRAINS_CFG self.scene.max_init_terrain_level = 0 + self.scene.terrain_curriculum.enabled = True + self.scene.terrain_curriculum.schedule = (0.0, 15000.0, 0.0, 9.0) self.scene.height_scanner.enable_height_scan = True self.scene.height_scanner.prim_body_name = "trunk" - self.scene.gemini2_camera.enable = True - self.scene.gemini2_camera.enable_rgb = False - self.scene.gemini2_camera.enable_depth = True - self.scene.gemini2_camera.partial_camera = True - self.scene.gemini2_camera.partial_camera_num_envs = 1024 - self.scene.gemini2_camera.partial_camera_seed = 42 - self.scene.gemini2_camera.partial_camera_force_tilt_crawl = True - self.scene.gemini2_camera.model_name = "wmp_front_depth" - self.scene.gemini2_camera.camera_model = "pinhole" - self.scene.gemini2_camera.spawn_prim_path = "trunk/wmp_depth_camera" - self.scene.gemini2_camera.spawn_offset_pos = (0.27, 0.0, 0.10) - self.scene.gemini2_camera.spawn_offset_rot = (1.0, 0.0, 0.0, 0.0) - self.scene.gemini2_camera.width = 64 - self.scene.gemini2_camera.height = 64 - self.scene.gemini2_camera.depth_near = 0.0 - self.scene.gemini2_camera.depth_far = 2.0 - self.scene.gemini2_camera.horizontal_fov_deg = 58.0 - self.scene.gemini2_camera.randomize_rotation = False - self.scene.gemini2_camera.random_pitch_deg = (0.0, 0.0) - self.scene.gemini2_camera.update_period = self.sim.dt * self.sim.decimation * wmp_update_interval - self.scene.gemini2_camera.allow_missing_depth_fallback = False + self.scene.rgbd_camera.enable = True + self.scene.rgbd_camera.enable_rgb = False + self.scene.rgbd_camera.enable_depth = True + self.scene.rgbd_camera.partial_camera = True + self.scene.rgbd_camera.partial_camera_num_envs = 1024 + self.scene.rgbd_camera.partial_camera_seed = 42 + self.scene.rgbd_camera.partial_camera_force_tilt_crawl = True + self.scene.rgbd_camera.model_name = "wmp_front_depth" + self.scene.rgbd_camera.camera_model = "pinhole" + self.scene.rgbd_camera.spawn_prim_path = "trunk/rgbd_camera" + self.scene.rgbd_camera.spawn_offset_pos = (0.3, 0.0, 0.03) + self.scene.rgbd_camera.spawn_offset_rot = (1.0, 0.0, 0.0, 0.0) + self.scene.rgbd_camera.width = 64 + self.scene.rgbd_camera.height = 64 + self.scene.rgbd_camera.depth_near = 0.0 + self.scene.rgbd_camera.depth_far = 2.0 + self.scene.rgbd_camera.horizontal_fov_deg = 58.0 + self.scene.rgbd_camera.randomize_rotation = True + self.scene.rgbd_camera.random_pitch_deg = (5.0, 15.0) + self.scene.rgbd_camera.update_period = self.sim.dt * self.sim.decimation * wmp_update_interval + self.scene.rgbd_camera.allow_missing_depth_fallback = False self.sim.render_interval = self.sim.decimation * wmp_update_interval self.robot.actor_obs_history_length = 5 self.robot.critic_obs_history_length = 1 + self.robot.wmp_privileged_contact_body_names = [".*thigh.*", ".*calf.*"] self.normalization.obs_scales.ang_vel = 0.25 self.normalization.obs_scales.joint_pos = 1.0 self.normalization.obs_scales.joint_vel = 0.05 self.normalization.obs_scales.height_scan = 5.0 + self.normalization.obs_scales.contact_force = 0.005 + self.normalization.obs_scales.pd_gains = 5.0 + self.normalization.obs_scales.com_pos = 20.0 self.normalization.clip_actions = 6.0 self.reward_settings.only_positive_rewards = True self.reward_settings.reward_curriculum = True @@ -244,7 +250,7 @@ def __post_init__(self): self.commands.ranges.ang_vel_z = (-1.0, 1.0) self.commands.ranges.heading = (0.0, 0.0) self.domain_rand.events.physics_material = EventTerm( - func=mdp.randomize_rigid_body_material, + func=mdp.wmp_recording_randomize_rigid_body_material, mode="startup", params={ "asset_cfg": SceneEntityCfg("robot", body_names=".*"), @@ -255,7 +261,7 @@ def __post_init__(self): }, ) self.domain_rand.events.add_base_mass = EventTerm( - func=mdp.randomize_rigid_body_mass, + func=mdp.wmp_recording_randomize_rigid_body_mass, mode="startup", params={ "asset_cfg": SceneEntityCfg("robot", body_names=".*trunk.*"), @@ -275,7 +281,7 @@ def __post_init__(self): }, ) self.domain_rand.events.randomize_base_com = EventTerm( - func=mdp.randomize_rigid_body_com, + func=mdp.wmp_recording_randomize_rigid_body_com, mode="startup", params={ "asset_cfg": SceneEntityCfg("robot", body_names=".*trunk.*"), @@ -283,7 +289,7 @@ def __post_init__(self): }, ) self.domain_rand.events.randomize_actuator_gains = EventTerm( - func=mdp.randomize_actuator_gains, + func=mdp.wmp_recording_randomize_actuator_gains, mode="startup", params={ "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), @@ -322,6 +328,7 @@ class A1WMPAMPTerrainAgentCfg(A1AMPFlatAgentCfg): "train_steps_per_iter": 10, "batch_size": 16, "batch_length": 64, + "model_lr": 1.0e-4, "replay_capacity_episodes": 50000, "replay_device": "cpu", "use_depth_predictor": True, @@ -366,11 +373,15 @@ def __post_init__(self): self.policy.actor_hidden_dims = [256, 128, 64] self.policy.critic_hidden_dims = [512, 256, 128] self.algorithm.class_name = "legged_lab.algorithms.wmp_amp_ppo:WMPAMPPPO" + self.algorithm.learning_rate = 3.0e-4 + self.algorithm.schedule = "fixed" + self.algorithm.desired_kl = 0.006 self.algorithm.entropy_coef = 0.01 self.algorithm.num_learning_epochs = 5 self.algorithm.num_mini_batches = 4 + self.algorithm.normalize_advantage_per_mini_batch = True self.algorithm.vel_predict_coef = 1.0 - self.algorithm.vel_target_start = 45 + self.algorithm.vel_target_start = 50 self.algorithm.vel_target_dim = 3 self.num_steps_per_env = 24 self.save_interval = 1000 diff --git a/legged_lab/envs/base/base_config.py b/legged_lab/envs/base/base_config.py index ddc7da2..a9fcae7 100644 --- a/legged_lab/envs/base/base_config.py +++ b/legged_lab/envs/base/base_config.py @@ -40,6 +40,14 @@ class RewardSettingsCfg: reward_curriculum_schedule: tuple[tuple[float, float, float, float], ...] = () +@configclass +class TerrainCurriculumCfg: + """按训练迭代逐步开放地形最大难度行。""" + + enabled: bool = False + schedule: tuple[float, float, float, float] = (0.0, 10000.0, 0.0, 5.0) + + @configclass class HeightScannerCfg: enable_height_scan: bool = False @@ -48,12 +56,13 @@ class HeightScannerCfg: size: tuple = (1.6, 1.0) forward_resolution: float = 0.1 forward_size: tuple = (2.0, 2.4) + forward_offset: tuple = (1.0, 0.0, 0.0) debug_vis: bool = False drift_range: tuple = (0.0, 0.0) @configclass -class Gemini2CameraCfg: +class RGBDCameraCfg: """RGBD 相机配置。 RGBD 相机统一在 InteractiveScene clone 完成后生成到机器人 body 下。 @@ -78,17 +87,19 @@ class Gemini2CameraCfg: partial_camera_num_envs: int | None = 1024 partial_camera_seed: int = 42 partial_camera_force_tilt_crawl: bool = True - spawn_prim_path: str = "base/wmp_depth_camera" + spawn_prim_path: str = "base/rgbd_camera" spawn_offset_pos: tuple[float, float, float] = (0.27, 0.0, 0.03) spawn_offset_rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) spawn_offset_convention: str = "world" horizontal_aperture: float = 20.955 vertical_aperture: float | None = None horizontal_fov_deg: float | None = 58.0 + show_visual_model: bool = False + visual_model_scale: float = 1.0 randomize_rotation: bool = True randomize_rotation_seed: int | None = None random_roll_deg: tuple[float, float] = (0.0, 0.0) - random_pitch_deg: tuple[float, float] = (-5.0, 5.0) + random_pitch_deg: tuple[float, float] = (5.0, 15.0) random_yaw_deg: tuple[float, float] = (0.0, 0.0) focal_length: float = 18.9002 focus_distance: float = 400.0 @@ -103,8 +114,9 @@ class BaseSceneCfg: terrain_type: str = MISSING terrain_generator: TerrainGeneratorCfg = None max_init_terrain_level: int = 5 + terrain_curriculum: TerrainCurriculumCfg = TerrainCurriculumCfg() height_scanner: HeightScannerCfg = HeightScannerCfg() - gemini2_camera: Gemini2CameraCfg = Gemini2CameraCfg() + rgbd_camera: RGBDCameraCfg = RGBDCameraCfg() @configclass @@ -115,6 +127,7 @@ class RobotCfg: wheel_velocity_scale: float = 8.0 terminate_contacts_body_names: list = [] feet_body_names: list = [] + wmp_privileged_contact_body_names: list = [] terminate_on_flight: bool = False terminate_on_flight_threshold: float = 1.0 @@ -129,6 +142,9 @@ class ObsScalesCfg: joint_vel: float = 1.0 actions: float = 1.0 height_scan: float = 1.0 + com_pos: float = 1.0 + pd_gains: float = 1.0 + contact_force: float = 0.01 @configclass diff --git a/legged_lab/envs/base/base_env.py b/legged_lab/envs/base/base_env.py index 2636775..0074f29 100644 --- a/legged_lab/envs/base/base_env.py +++ b/legged_lab/envs/base/base_env.py @@ -14,6 +14,7 @@ import math import numpy as np import torch +from torch.nn import functional as F from isaaclab.assets.articulation import Articulation from isaaclab.envs.mdp.commands import UniformVelocityCommand, UniformVelocityCommandCfg from isaaclab.managers import EventManager, RewardManager @@ -21,13 +22,13 @@ from isaaclab.scene import InteractiveScene from isaaclab.sensors import ContactSensor, RayCaster, TiledCameraCfg from isaaclab.sim import PhysxCfg, SimulationContext -from isaaclab.utils.math import quat_from_euler_xyz, quat_mul +from isaaclab.utils.math import quat_from_angle_axis, quat_mul from isaaclab.utils.buffers import CircularBuffer, DelayBuffer from rsl_rl.env import VecEnv from tensordict import TensorDict from legged_lab.envs.base.base_env_config import BaseEnvCfg -from legged_lab.sensors import WMPPartialTiledCamera, WMPPartialTiledCameraCfg, select_wmp_camera_env_ids +from legged_lab.sensors import WMPPartialTiledCamera, WMPPartialTiledCameraCfg, select_wmp_depth_indices from legged_lab.utils.env_utils.scene import SceneCfg @@ -38,6 +39,10 @@ def __init__(self, cfg: BaseEnvCfg, headless): self.cfg = cfg self.headless = headless self.device = self.cfg.device + self.depth_index_without_crawl_tilt = torch.empty(0, device=self.device, dtype=torch.long) + self.depth_index = torch.empty(0, device=self.device, dtype=torch.long) + self.depth_index_inverse = torch.empty(0, device=self.device, dtype=torch.long) + self.wmp_depth_buffer = None self.physics_dt = self.cfg.sim.dt self.step_dt = self.cfg.sim.decimation * self.cfg.sim.dt render_interval = self.cfg.sim.render_interval @@ -76,6 +81,8 @@ def __init__(self, cfg: BaseEnvCfg, headless): self.x_edge_mask = getattr(self.scene.terrain, "x_edge_mask", None) self.wmp_edge_query_offset = getattr(self.scene.terrain, "wmp_edge_query_offset", (0.0, 0.0)) self.wmp_terrain_horizontal_scale = getattr(self.scene.terrain, "wmp_horizontal_scale", 1.0) + self.wmp_terrain_col_kinds = getattr(self.scene.terrain, "wmp_terrain_col_kinds", ()) + self.wmp_terrain_cols_by_kind = getattr(self.scene.terrain, "wmp_terrain_cols_by_kind", {}) self.terrain_levels = getattr(self.scene.terrain, "terrain_levels", None) self.terrain_types = getattr(self.scene.terrain, "terrain_types", None) self.gap_start_col = getattr(self.scene.terrain, "gap_start_col", 0) @@ -85,21 +92,35 @@ def __init__(self, cfg: BaseEnvCfg, headless): "[INFO] WMP x_edge_mask enabled: " f"shape={tuple(self.x_edge_mask.shape)}, true_count={int(self.x_edge_mask.sum().item())}, " f"offset={self.wmp_edge_query_offset}, horizontal_scale={self.wmp_terrain_horizontal_scale}, " - f"gap/climb cols=[{self.gap_start_col}, {self.climb_end_col})" + f"gap/climb cols=[{self.gap_start_col}, {self.climb_end_col}), " + f"col_kinds={self.wmp_terrain_col_kinds}" ) if self.cfg.scene.height_scanner.enable_height_scan: self.height_scanner: RayCaster = self.scene.sensors["height_scanner"] self.forward_height_scanner: RayCaster = self.scene.sensors["forward_height_scanner"] - self.gemini2_depth_camera = self.scene.sensors.get("gemini2_depth_camera") - if self.gemini2_depth_camera is not None and hasattr(self.gemini2_depth_camera, "camera_env_ids"): - camera_ids = self.gemini2_depth_camera.camera_env_ids.to(self.device) + self.rgbd_camera = self.scene.sensors.get("rgbd_camera") + if self.rgbd_camera is not None and hasattr(self.rgbd_camera, "camera_env_ids"): + camera_ids = self.rgbd_camera.camera_env_ids.to(self.device) self.wmp_camera_env_ids = camera_ids self.wmp_camera_env_mask = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) self.wmp_camera_env_mask[camera_ids] = True + self.depth_index = camera_ids + if hasattr(self.rgbd_camera, "camera_env_id_inverse"): + self.depth_index_inverse = self.rgbd_camera.camera_env_id_inverse.to(self.device) + else: + self.depth_index_inverse = torch.full((self.num_envs,), -1, device=self.device, dtype=torch.long) + self.depth_index_inverse[self.depth_index] = torch.arange(self.depth_index.numel(), device=self.device) + if self.depth_index_without_crawl_tilt.numel() > 0: + valid = self.depth_index_inverse[self.depth_index_without_crawl_tilt] >= 0 + self.depth_index_without_crawl_tilt = self.depth_index_without_crawl_tilt[valid] + self._init_wmp_depth_buffer() print( - "[INFO] WMP partial depth camera enabled: " - f"model={getattr(self.gemini2_depth_camera.cfg, 'camera_model_name', 'unknown')}, " - f"camera_envs={int(camera_ids.numel())}/{self.num_envs}" + "[INFO] WMP partial RGBD camera enabled: " + f"model={getattr(self.rgbd_camera.cfg, 'camera_model_name', 'unknown')}, " + f"camera_envs={int(camera_ids.numel())}/{self.num_envs}, " + f"depth_index_count={int(self.depth_index.numel())}, " + f"depth_index_without_crawl_tilt_count={int(self.depth_index_without_crawl_tilt.numel())}, " + f"tilt_crawl_depth_count={int(self.depth_index.numel() - self.depth_index_without_crawl_tilt.numel())}" ) command_cfg = UniformVelocityCommandCfg( @@ -115,32 +136,40 @@ def __init__(self, cfg: BaseEnvCfg, headless): self.command_generator = UniformVelocityCommand(cfg=command_cfg, env=self) self.reward_manager = RewardManager(self.cfg.reward, self) self.reward_curriculum_coef = self._init_reward_curriculum_coef() + self.terrain_curriculum_max_level = self._terrain_curriculum_allowed_level(0) self.init_buffers() env_ids = torch.arange(self.num_envs, device=self.device) self.event_manager = EventManager(self.cfg.domain_rand.events, self) if "startup" in self.event_manager.available_modes: self.event_manager.apply(mode="startup") + self._refresh_wmp_privileged_buffers_from_sim() self.reset(env_ids) self._is_closed = False def _attach_rgbd_cameras_after_scene_creation(self): - camera = self.cfg.scene.gemini2_camera + camera = self.cfg.scene.rgbd_camera if not camera.enable: return camera_env_ids = () camera_cfg_cls = TiledCameraCfg if camera.partial_camera: - camera_env_ids_tensor = select_wmp_camera_env_ids( + terrain = getattr(self.scene, "terrain", None) + depth_index_without_crawl_tilt, depth_index, depth_index_inverse = select_wmp_depth_indices( num_envs=self.cfg.scene.num_envs, camera_num_envs=camera.partial_camera_num_envs, seed=camera.partial_camera_seed, terrain_generator=self.cfg.scene.terrain_generator, + terrain_types=getattr(terrain, "terrain_types", None), + terrain_cols_by_kind=getattr(terrain, "wmp_terrain_cols_by_kind", None), force_tilt_crawl=camera.partial_camera_force_tilt_crawl, device="cpu", ) - camera_env_ids = tuple(int(env_id) for env_id in camera_env_ids_tensor.tolist()) + self.depth_index_without_crawl_tilt = depth_index_without_crawl_tilt.to(self.device) + self.depth_index = depth_index.to(self.device) + self.depth_index_inverse = depth_index_inverse.to(self.device) + camera_env_ids = tuple(int(env_id) for env_id in depth_index.tolist()) camera_cfg_cls = WMPPartialTiledCameraCfg prim_path = "{ENV_REGEX_NS}/Robot/" + camera.spawn_prim_path.strip("/") @@ -164,13 +193,44 @@ def _attach_rgbd_cameras_after_scene_creation(self): rgbd_camera = self._make_camera_cfg( camera_cfg_cls, prim_path, offset_cfg, spawn_cfg, data_types, camera_env_ids, camera_rot ) - if camera.enable_depth: - self.scene._sensors["gemini2_depth_camera"] = rgbd_camera - elif camera.enable_rgb: - self.scene._sensors["gemini2_rgb_camera"] = rgbd_camera + if camera.show_visual_model: + self._spawn_camera_visual_model(prim_path, camera_env_ids) + self.scene._sensors["rgbd_camera"] = rgbd_camera + + def _spawn_camera_visual_model(self, prim_path: str, camera_env_ids): + """在真实 USD Camera 下挂一个纯视觉外壳，便于播放时确认相机安装位置。""" + camera = self.cfg.scene.rgbd_camera + stage = sim_utils.get_current_stage() + if camera_env_ids: + camera_paths = [prim_path.replace("env_.*", f"env_{int(env_id)}", 1) for env_id in camera_env_ids] + else: + camera_paths = [prim_path.replace("env_.*", f"env_{env_id}", 1) for env_id in range(self.num_envs)] + + scale = float(camera.visual_model_scale) + body_cfg = sim_utils.CuboidCfg( + size=(0.055 * scale, 0.040 * scale, 0.032 * scale), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.02, 0.02, 0.025), roughness=0.65), + ) + lens_cfg = sim_utils.CylinderCfg( + radius=0.013 * scale, + height=0.018 * scale, + axis="Z", + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.02, 0.12, 0.45), roughness=0.25), + ) + # USD Camera 局部坐标使用 OpenGL 约定: -Z 为朝向，+Y 向上。 + # 这些可视几何放在相机后方(+Z)，避免挡住真实深度相机的视锥。 + for path in camera_paths: + if not stage.GetPrimAtPath(path).IsValid(): + continue + body_path = f"{path}/visual_body" + lens_path = f"{path}/visual_lens" + if not stage.GetPrimAtPath(body_path).IsValid(): + body_cfg.func(body_path, body_cfg, translation=(0.0, 0.0, 0.035 * scale)) + if not stage.GetPrimAtPath(lens_path).IsValid(): + lens_cfg.func(lens_path, lens_cfg, translation=(0.0, 0.0, 0.005 * scale)) def _make_camera_spawn_cfg(self): - camera = self.cfg.scene.gemini2_camera + camera = self.cfg.scene.rgbd_camera focal_length = camera.focal_length if camera.horizontal_fov_deg is not None: # USD pinhole 相机水平视场: @@ -191,7 +251,7 @@ def _make_camera_spawn_cfg(self): raise ValueError(f"Unsupported camera_model: {camera.camera_model}") def _make_camera_offset_rotations(self, camera_env_ids): - camera = self.cfg.scene.gemini2_camera + camera = self.cfg.scene.rgbd_camera if not camera.randomize_rotation: return None count = len(camera_env_ids) if camera_env_ids else self.num_envs @@ -207,15 +267,19 @@ def sample_degrees(bounds): return torch.deg2rad(values * (high - low) + low) base = torch.tensor(camera.spawn_offset_rot, dtype=torch.float32).repeat(count, 1) - delta = quat_from_euler_xyz( - sample_degrees(camera.random_roll_deg), - sample_degrees(camera.random_pitch_deg), - sample_degrees(camera.random_yaw_deg), - ) - return quat_mul(base, delta) + roll = sample_degrees(camera.random_roll_deg) + pitch = sample_degrees(camera.random_pitch_deg) + yaw = sample_degrees(camera.random_yaw_deg) + local_axes = torch.eye(3, dtype=torch.float32) + q_roll = quat_from_angle_axis(roll, local_axes[0].repeat(count, 1)) + q_pitch = quat_from_angle_axis(pitch, local_axes[1].repeat(count, 1)) + q_yaw = quat_from_angle_axis(yaw, local_axes[2].repeat(count, 1)) + # world 相机约定: forward=+X, up=+Z。四元数作用为 v_w = q * v_cam * q^{-1}。 + # 因此 pitch 应绕相机局部 +Y 轴旋转；正角会把 +X 转向 -Z，即让相机向下俯视。 + return quat_mul(quat_mul(quat_mul(base, q_roll), q_pitch), q_yaw) def _make_camera_cfg(self, camera_cfg_cls, prim_path, offset_cfg, spawn_cfg, data_types, camera_env_ids, camera_rot): - camera = self.cfg.scene.gemini2_camera + camera = self.cfg.scene.rgbd_camera kwargs = dict( prim_path=prim_path, offset=offset_cfg, @@ -277,6 +341,12 @@ def init_buffers(self): self.termination_contact_cfg.resolve(self.scene) self.feet_cfg = SceneEntityCfg(name="contact_sensor", body_names=self.cfg.robot.feet_body_names) self.feet_cfg.resolve(self.scene) + self.wmp_privileged_contact_cfg = None + if self._use_original_wmp_critic_obs(): + self.wmp_privileged_contact_cfg = SceneEntityCfg( + name="contact_sensor", body_names=self.cfg.robot.wmp_privileged_contact_body_names + ) + self.wmp_privileged_contact_cfg.resolve(self.scene) self._amp_order_logged = False self.obs_scales = self.cfg.normalization.obs_scales @@ -285,6 +355,7 @@ def init_buffers(self): self.episode_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long) self.sim_step_counter = 0 self.time_out_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + self._init_wmp_privileged_buffers() self.init_obs_buffer() def compute_current_observations(self): @@ -310,10 +381,17 @@ def compute_current_observations(self): ) root_lin_vel = robot.data.root_lin_vel_b - feet_contact = torch.max(torch.norm(net_contact_forces[:, :, self.feet_cfg.body_ids], dim=-1), dim=1)[0] > 0.5 - current_critic_obs = torch.cat( - [current_actor_obs, root_lin_vel * self.obs_scales.lin_vel, feet_contact], dim=-1 - ) + if self._use_original_wmp_critic_obs(): + current_critic_obs = self._compute_original_wmp_critic_obs( + current_actor_obs=current_actor_obs, + root_lin_vel=root_lin_vel, + net_contact_forces=net_contact_forces, + ) + else: + feet_contact = torch.max(torch.norm(net_contact_forces[:, :, self.feet_cfg.body_ids], dim=-1), dim=1)[0] > 0.5 + current_critic_obs = torch.cat( + [current_actor_obs, root_lin_vel * self.obs_scales.lin_vel, feet_contact], dim=-1 + ) return current_actor_obs, current_critic_obs @@ -337,6 +415,8 @@ def compute_observations(self): if self.add_noise: height_scan += (2 * torch.rand_like(height_scan) - 1) * self.height_scan_noise_vec actor_obs = torch.cat([actor_obs, height_scan], dim=-1) + if self._use_original_wmp_critic_obs() and critic_obs.shape[-1] != 285: + raise RuntimeError(f"Original WMP critic obs must be 285 dim, got {critic_obs.shape[-1]}.") actor_obs = torch.clip(actor_obs, -self.clip_obs, self.clip_obs) critic_obs = torch.clip(critic_obs, -self.clip_obs, self.clip_obs) @@ -416,6 +496,7 @@ def step(self, actions: torch.Tensor): self.command_generator.compute(self.step_dt) if "interval" in self.event_manager.available_modes: self.event_manager.apply(mode="interval", dt=self.step_dt) + self._update_wmp_depth_buffer() self.reset_buf, self.time_out_buf = self.check_reset() reward_buf = self.reward_manager.compute(self.step_dt) @@ -433,21 +514,102 @@ def step(self, actions: torch.Tensor): return obs, reward_buf, self.reset_buf, self.extras def get_depth_observations(self): - if self.gemini2_depth_camera is None: - if self.cfg.scene.gemini2_camera.allow_missing_depth_fallback: + if self.rgbd_camera is None: + if self.cfg.scene.rgbd_camera.allow_missing_depth_fallback: return torch.full( - (self.num_envs, self.cfg.scene.gemini2_camera.height, self.cfg.scene.gemini2_camera.width, 1), - self.cfg.scene.gemini2_camera.depth_far, + (self.num_envs, self.cfg.scene.rgbd_camera.height, self.cfg.scene.rgbd_camera.width, 1), + self.cfg.scene.rgbd_camera.depth_far, device=self.device, ) - raise RuntimeError("Gemini2 depth camera is not enabled for this task.") - return self.gemini2_depth_camera.data.output["distance_to_image_plane"] + raise RuntimeError("RGBD camera is not enabled for this task.") + return self.rgbd_camera.data.output["distance_to_image_plane"] + + def _init_wmp_depth_buffer(self): + camera = self.cfg.scene.rgbd_camera + if self.depth_index.numel() == 0 or not camera.enable_depth: + self.wmp_depth_buffer = None + return + self.wmp_depth_buffer = torch.zeros( + (self.depth_index.numel(), 2, camera.height, camera.width, 1), + device=self.device, + dtype=torch.float32, + ) + + def _update_wmp_depth_buffer(self): + if self.wmp_depth_buffer is None or self.rgbd_camera is None: + return + if self.sim_step_counter % self.render_interval != 0: + return + depth = self.rgbd_camera.data.output.get("distance_to_image_plane") + if depth is None: + return + image = self._depth_to_wmp_image(depth.to(self.device)) + if image.shape[0] != self.depth_index.numel(): + raise RuntimeError( + "WMP depth buffer expects partial camera batch to match depth_index: " + f"depth_batch={image.shape[0]}, depth_index={self.depth_index.numel()}" + ) + reset_like = self.episode_length_buf[self.depth_index] <= 1 + self.wmp_depth_buffer[:, 0].copy_(self.wmp_depth_buffer[:, 1]) + self.wmp_depth_buffer[:, 1].copy_(image) + if torch.any(reset_like): + ids = reset_like.nonzero(as_tuple=False).flatten() + self.wmp_depth_buffer[ids, 0].copy_(image[ids]) + self.wmp_depth_buffer[ids, 1].copy_(image[ids]) + + def _depth_to_wmp_image(self, depth: torch.Tensor) -> torch.Tensor: + camera = self.cfg.scene.rgbd_camera + near = float(camera.depth_near) + far = float(camera.depth_far) + depth = torch.nan_to_num(depth.float(), nan=far, posinf=far, neginf=-far) + if depth.ndim != 4: + raise ValueError(f"depth must have shape B,H,W,1 or B,1,H,W, got {tuple(depth.shape)}") + if depth.shape[-1] == 1: + depth = depth.permute(0, 3, 1, 2) + elif depth.shape[1] != 1: + raise ValueError(f"depth channel dimension must be 1, got {tuple(depth.shape)}") + # 原版 IsaacGym 为负深度: image = ((-depth) - near) / (far - near) - 0.5。 + # IsaacLab distance_to_image_plane 为正深度；若后端给负数，则先取反统一成正距离。 + depth = torch.where(depth < 0.0, -depth, depth) + depth = torch.clamp(depth, near, far) + depth = (depth - near) / max(far - near, 1.0e-6) - 0.5 + image = F.interpolate( + depth, + size=(camera.height, camera.width), + mode="bilinear", + align_corners=False, + ) + return image.permute(0, 2, 3, 1).contiguous() + + def get_wmp_depth_observations(self): + if self.wmp_depth_buffer is None: + raise RuntimeError("WMP depth buffer is not initialized.") + # 原版返回 depth_buffer[:, -2]；最后一帧已经写入 buffer，但 WMP 消费倒数第二帧。 + return self.wmp_depth_buffer[:, 0] def get_depth_camera_env_ids(self): - if self.gemini2_depth_camera is not None and hasattr(self.gemini2_depth_camera, "camera_env_ids"): - return self.gemini2_depth_camera.camera_env_ids.to(self.device) + if self.depth_index.numel() > 0: + return self.depth_index.to(self.device) + if self.rgbd_camera is not None and hasattr(self.rgbd_camera, "camera_env_ids"): + return self.rgbd_camera.camera_env_ids.to(self.device) return torch.arange(self.num_envs, device=self.device, dtype=torch.long) + def get_depth_index(self): + if self.depth_index.numel() > 0: + return self.depth_index.to(self.device) + return self.get_depth_camera_env_ids() + + def get_depth_index_without_crawl_tilt(self): + return self.depth_index_without_crawl_tilt.to(self.device) + + def get_depth_index_inverse(self): + if self.depth_index_inverse.numel() == self.num_envs: + return self.depth_index_inverse.to(self.device) + inverse = torch.full((self.num_envs,), -1, device=self.device, dtype=torch.long) + depth_index = self.get_depth_index() + inverse[depth_index] = torch.arange(depth_index.numel(), device=self.device, dtype=torch.long) + return inverse + def get_wmp_proprioception(self): robot = self.robot command = self.command_generator.command[:, :3] @@ -553,6 +715,180 @@ def init_obs_buffer(self): max_len=self.cfg.robot.critic_obs_history_length, batch_size=self.num_envs, device=self.device ) + def _init_wmp_privileged_buffers(self): + """缓存原版 WMP critic 需要的 privileged 域随机化量。 + + 原版 critic 布局: + contact_flag(8) + contact_force(12) + d_gain_scale(12) + p_gain_scale(12) + + com_pos(3) + added_mass(1) + restitution(1) + friction(1) + + base_lin_vel(3) + actor_obs(45) = 98 + + 再在 `compute_observations()` 末尾拼接 height_scan(187)，得到原版的 285 维 critic obs。 + A1 WMP 使用记录型 startup EventTerm，同一次 sample 会同时写入物理仿真和这些 buffer； + 这里的初始化值只作为 EventTerm 缺失时的保底。 + """ + if not self._use_original_wmp_critic_obs(): + return + self.wmp_priv_friction = torch.zeros(self.num_envs, 1, device=self.device) + self.wmp_priv_restitution = torch.zeros(self.num_envs, 1, device=self.device) + self.wmp_priv_added_mass = torch.zeros(self.num_envs, 1, device=self.device) + self.wmp_priv_com_pos = torch.zeros(self.num_envs, 3, device=self.device) + self.wmp_priv_p_gain_scale = torch.zeros(self.num_envs, self.num_actions, device=self.device) + self.wmp_priv_d_gain_scale = torch.zeros(self.num_envs, self.num_actions, device=self.device) + self.wmp_priv_default_coms = self.robot.root_physx_view.get_coms().to(self.device).clone() + self.wmp_critic_obs_slices = { + "contact_flag": slice(0, 8), + "contact_force": slice(8, 20), + "d_gain_scale": slice(20, 32), + "p_gain_scale": slice(32, 44), + "com_pos": slice(44, 47), + "added_mass": slice(47, 48), + "restitution": slice(48, 49), + "friction": slice(49, 50), + "base_lin_vel": slice(50, 53), + "actor_obs": slice(53, 98), + "height_scan": slice(98, 285), + } + self._sample_wmp_privileged_randomization(torch.arange(self.num_envs, device=self.device)) + + def _use_original_wmp_critic_obs(self) -> bool: + return len(self.cfg.robot.wmp_privileged_contact_body_names) > 0 + + def _sample_wmp_privileged_randomization(self, env_ids: torch.Tensor): + if len(env_ids) == 0: + return + + def uniform(bounds, shape): + low, high = float(bounds[0]), float(bounds[1]) + if abs(high - low) < 1.0e-8: + return torch.full(shape, low, device=self.device) + return torch.empty(shape, device=self.device).uniform_(low, high) + + events = self.cfg.domain_rand.events + physics_material = getattr(events, "physics_material", None) + if physics_material is not None: + params = physics_material.params + self.wmp_priv_friction[env_ids] = uniform(params.get("static_friction_range", (1.0, 1.0)), (len(env_ids), 1)) + self.wmp_priv_restitution[env_ids] = uniform(params.get("restitution_range", (0.0, 0.0)), (len(env_ids), 1)) + + add_base_mass = getattr(events, "add_base_mass", None) + if add_base_mass is not None: + params = add_base_mass.params + if params.get("operation", "add") == "add": + self.wmp_priv_added_mass[env_ids] = uniform( + params.get("mass_distribution_params", (0.0, 0.0)), (len(env_ids), 1) + ) + + randomize_base_com = getattr(events, "randomize_base_com", None) + if randomize_base_com is not None: + com_range = randomize_base_com.params.get("com_range", {}) + self.wmp_priv_com_pos[env_ids, 0] = uniform(com_range.get("x", (0.0, 0.0)), (len(env_ids),)) + self.wmp_priv_com_pos[env_ids, 1] = uniform(com_range.get("y", (0.0, 0.0)), (len(env_ids),)) + self.wmp_priv_com_pos[env_ids, 2] = uniform(com_range.get("z", (0.0, 0.0)), (len(env_ids),)) + + randomize_actuator_gains = getattr(events, "randomize_actuator_gains", None) + if randomize_actuator_gains is not None: + params = randomize_actuator_gains.params + if params.get("operation", "scale") == "scale": + self.wmp_priv_p_gain_scale[env_ids] = uniform( + params.get("stiffness_distribution_params", (1.0, 1.0)), (len(env_ids), self.num_actions) + ) - 1.0 + self.wmp_priv_d_gain_scale[env_ids] = uniform( + params.get("damping_distribution_params", (1.0, 1.0)), (len(env_ids), self.num_actions) + ) - 1.0 + + def _refresh_wmp_privileged_buffers_from_sim(self): + if not self._use_original_wmp_critic_obs(): + return + env_ids = torch.arange(self.num_envs, device=self.device) + + try: + materials = self.robot.root_physx_view.get_material_properties().to(self.device) + if materials.ndim == 3 and materials.shape[-1] >= 3: + self.wmp_priv_friction[env_ids] = materials[..., 0].mean(dim=1, keepdim=True) + self.wmp_priv_restitution[env_ids] = materials[..., 2].mean(dim=1, keepdim=True) + except Exception as exc: + print(f"[WARN] Failed to sync WMP material privileged obs from sim: {exc}") + + try: + body_id = self._first_event_body_id("add_base_mass") + masses = self.robot.root_physx_view.get_masses().to(self.device) + default_masses = self.robot.data.default_mass.to(self.device) + self.wmp_priv_added_mass[env_ids] = masses[:, body_id : body_id + 1] - default_masses[:, body_id : body_id + 1] + except Exception as exc: + print(f"[WARN] Failed to sync WMP mass privileged obs from sim: {exc}") + + try: + body_id = self._first_event_body_id("randomize_base_com") + coms = self.robot.root_physx_view.get_coms().to(self.device) + self.wmp_priv_com_pos[env_ids] = coms[:, body_id, :3] - self.wmp_priv_default_coms[:, body_id, :3] + except Exception as exc: + print(f"[WARN] Failed to sync WMP COM privileged obs from sim: {exc}") + + try: + default_stiffness = self.robot.data.default_joint_stiffness.to(self.device) + default_damping = self.robot.data.default_joint_damping.to(self.device) + stiffness = self.robot.data.joint_stiffness.to(self.device) + damping = self.robot.data.joint_damping.to(self.device) + self.wmp_priv_p_gain_scale[env_ids] = stiffness / torch.clamp(default_stiffness, min=1.0e-8) - 1.0 + self.wmp_priv_d_gain_scale[env_ids] = damping / torch.clamp(default_damping, min=1.0e-8) - 1.0 + except Exception as exc: + print(f"[WARN] Failed to sync WMP gain privileged obs from sim: {exc}") + + def _first_event_body_id(self, event_name: str) -> int: + event = getattr(self.cfg.domain_rand.events, event_name, None) + asset_cfg = None if event is None else event.params.get("asset_cfg") + body_ids = getattr(asset_cfg, "body_ids", None) + if isinstance(body_ids, slice) or body_ids is None: + return 0 + if isinstance(body_ids, torch.Tensor): + return int(body_ids.flatten()[0].item()) if body_ids.numel() > 0 else 0 + return int(body_ids[0]) if len(body_ids) > 0 else 0 + + def _fit_last_dim(self, value: torch.Tensor, target_dim: int) -> torch.Tensor: + dim = value.shape[-1] + if dim == target_dim: + return value + if dim > target_dim: + return value[..., :target_dim] + pad_shape = (*value.shape[:-1], target_dim - dim) + return torch.cat([value, torch.zeros(pad_shape, device=value.device, dtype=value.dtype)], dim=-1) + + def _compute_original_wmp_critic_obs( + self, + current_actor_obs: torch.Tensor, + root_lin_vel: torch.Tensor, + net_contact_forces: torch.Tensor, + ) -> torch.Tensor: + latest_forces = net_contact_forces[:, -1] + contact_force = latest_forces[:, self.feet_cfg.body_ids].reshape(self.num_envs, -1) + contact_force = self._fit_last_dim(contact_force, 12) * self.obs_scales.contact_force + + if self.wmp_privileged_contact_cfg is not None: + contact_flag = torch.norm(latest_forces[:, self.wmp_privileged_contact_cfg.body_ids], dim=-1) > 0.1 + else: + contact_flag = torch.zeros(self.num_envs, 0, device=self.device, dtype=torch.bool) + contact_flag = self._fit_last_dim(contact_flag.float(), 8) + + current_critic_obs = torch.cat( + [ + contact_flag, + contact_force, + self.wmp_priv_d_gain_scale * self.obs_scales.pd_gains, + self.wmp_priv_p_gain_scale * self.obs_scales.pd_gains, + self.wmp_priv_com_pos * self.obs_scales.com_pos, + self.wmp_priv_added_mass, + self.wmp_priv_restitution, + self.wmp_priv_friction, + root_lin_vel * self.obs_scales.lin_vel, + current_actor_obs, + ], + dim=-1, + ) + if current_critic_obs.shape[-1] != 98: + raise RuntimeError(f"Original WMP critic current obs must be 98 dim, got {current_critic_obs.shape[-1]}.") + return current_critic_obs + def update_terrain_levels(self, env_ids): distance = torch.norm(self.robot.data.root_pos_w[env_ids, :2] - self.scene.env_origins[env_ids, :2], dim=1) move_up = distance > self.scene.terrain.cfg.terrain_generator.size[0] / 2 @@ -561,9 +897,46 @@ def update_terrain_levels(self, env_ids): ) move_down *= ~move_up self.scene.terrain.update_env_origins(env_ids, move_up, move_down) + self._clamp_terrain_levels(env_ids) extras = {"Curriculum/terrain_levels": torch.mean(self.scene.terrain.terrain_levels.float())} + if self.terrain_curriculum_max_level is not None: + extras["Curriculum/terrain_max_allowed_level"] = float(self.terrain_curriculum_max_level) return extras + def _terrain_curriculum_allowed_level(self, current_iter: int) -> int | None: + settings = getattr(self.cfg.scene, "terrain_curriculum", None) + if settings is None or not settings.enabled: + return None + terrain = getattr(self.scene, "terrain", None) if hasattr(self, "scene") else None + max_level = int(getattr(terrain, "max_terrain_level", 1)) - 1 if terrain is not None else 0 + start_iter, end_iter, start_level, end_level = [float(x) for x in settings.schedule] + # 线性课程公式： + # alpha = clamp((iter - start_iter) / (end_iter - start_iter), 0, 1) + # level = round((1 - alpha) * start_level + alpha * end_level) + alpha = 1.0 if abs(end_iter - start_iter) < 1.0e-8 else (float(current_iter) - start_iter) / (end_iter - start_iter) + alpha = max(0.0, min(1.0, alpha)) + level = round((1.0 - alpha) * start_level + alpha * end_level) + return max(0, min(int(level), max_level)) + + def _clamp_terrain_levels(self, env_ids: torch.Tensor): + max_level = self.terrain_curriculum_max_level + terrain = getattr(self.scene, "terrain", None) + if max_level is None or terrain is None or getattr(terrain, "terrain_origins", None) is None: + return + terrain.terrain_levels[env_ids] = torch.clamp(terrain.terrain_levels[env_ids], max=int(max_level)) + terrain.env_origins[env_ids] = terrain.terrain_origins[terrain.terrain_levels[env_ids], terrain.terrain_types[env_ids]] + + def update_training_curriculum(self, current_iter: int) -> dict[str, float]: + logs = self.update_reward_curriculum(current_iter) + allowed_level = self._terrain_curriculum_allowed_level(current_iter) + if allowed_level is None: + return logs + self.terrain_curriculum_max_level = allowed_level + env_ids = torch.arange(self.num_envs, device=self.device) + self._clamp_terrain_levels(env_ids) + logs["Curriculum/terrain_max_allowed_level"] = float(allowed_level) + return logs + def _init_reward_curriculum_coef(self) -> dict[str, float]: settings = self.cfg.reward_settings terms = tuple(settings.reward_curriculum_term) @@ -642,7 +1015,7 @@ def close(self): print("[INFO] Closing BaseEnv resources.") if hasattr(self, "scene"): sensors = getattr(self.scene, "sensors", {}) - for name in ("gemini2_rgb_camera", "gemini2_depth_camera"): + for name in ("rgbd_camera",): if name in sensors: print(f"[INFO] Releasing sensor: {name}") del self.scene diff --git a/legged_lab/scripts/convert_original_wmp_checkpoint.py b/legged_lab/scripts/convert_original_wmp_checkpoint.py new file mode 100644 index 0000000..a5045d6 --- /dev/null +++ b/legged_lab/scripts/convert_original_wmp_checkpoint.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""将 ByteDance 原版 WMP checkpoint 转成当前 LeggedLab WMP 播放格式。""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch + + +def _rename_world_model_key(key: str) -> str: + replacements = ( + ("encoder._mlp.layers.Encoder_linear", "encoder._mlp.layers.linear"), + ("encoder._mlp.layers.Encoder_norm", "encoder._mlp.layers.norm"), + ("heads.decoder._mlp.layers.Decoder_linear", "heads.decoder._mlp.layers.linear"), + ("heads.decoder._mlp.layers.Decoder_norm", "heads.decoder._mlp.layers.norm"), + ("heads.reward.layers.Reward_linear", "heads.reward.layers.linear"), + ("heads.reward.layers.Reward_norm", "heads.reward.layers.norm"), + ("dynamics._cell.layers.GRU_linear.weight", "dynamics._cell.layers.0.weight"), + ) + for old, new in replacements: + key = key.replace(old, new) + return key + + +def _convert_world_model(src: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return {_rename_world_model_key(key): value for key, value in src.items()} + + +def _convert_actor(src: dict[str, torch.Tensor], reference: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + dst = {key: value.clone() for key, value in reference.items()} + mapping = { + "std": "distribution.std_param", + "actor.0.weight": "mlp.0.weight", + "actor.0.bias": "mlp.0.bias", + "actor.2.weight": "mlp.2.weight", + "actor.2.bias": "mlp.2.bias", + "actor.4.weight": "mlp.4.weight", + "actor.4.bias": "mlp.4.bias", + "actor.6.weight": "mlp.6.weight", + "actor.6.bias": "mlp.6.bias", + "wm_feature_encoder.0.weight": "wmp_encoder.0.weight", + "wm_feature_encoder.0.bias": "wmp_encoder.0.bias", + "wm_feature_encoder.2.weight": "wmp_encoder.2.weight", + "wm_feature_encoder.2.bias": "wmp_encoder.2.bias", + "wm_feature_encoder.4.weight": "wmp_encoder.4.weight", + "wm_feature_encoder.4.bias": "wmp_encoder.4.bias", + "history_encoder.0.weight": "history_encoder.0.weight", + "history_encoder.0.bias": "history_encoder.0.bias", + "history_encoder.2.weight": "history_encoder.2.weight", + "history_encoder.2.bias": "history_encoder.2.bias", + "history_encoder.4.weight": "history_encoder.4.weight", + "history_encoder.4.bias": "history_encoder.4.bias", + } + for src_key, dst_key in mapping.items(): + if src_key not in src or dst_key not in dst: + raise KeyError(f"Missing actor key mapping {src_key} -> {dst_key}") + if tuple(src[src_key].shape) != tuple(dst[dst_key].shape): + raise ValueError(f"Actor shape mismatch {src_key} -> {dst_key}: {src[src_key].shape} vs {dst[dst_key].shape}") + dst[dst_key] = src[src_key].clone() + # 原版 ActorCriticWMP 没有 actor latent normalizer；当前 WMPMLPModel 有 normalizer。 + # 为保持原版推理语义，设为恒等变换: norm(x) = (x - 0) / 1。 + for key in ("obs_normalizer._mean",): + if key in dst: + dst[key] = torch.zeros_like(dst[key]) + for key in ("obs_normalizer._var", "obs_normalizer._std"): + if key in dst: + dst[key] = torch.ones_like(dst[key]) + return dst + + +def _copy_matching(src: dict[str, torch.Tensor], reference: dict[str, torch.Tensor], prefix: str) -> dict[str, torch.Tensor]: + dst = {key: value.clone() for key, value in reference.items()} + for key, value in src.items(): + if key not in dst: + raise KeyError(f"{prefix}: unexpected key {key}") + if tuple(value.shape) != tuple(dst[key].shape): + raise ValueError(f"{prefix}: shape mismatch {key}: {value.shape} vs {dst[key].shape}") + dst[key] = value.clone() + return dst + + +def main() -> None: + parser = argparse.ArgumentParser(description="Convert original ByteDance WMP checkpoint to LeggedLab format.") + parser.add_argument("source", type=Path, help="Original WMP checkpoint path.") + parser.add_argument( + "--reference", + type=Path, + required=True, + help="A LeggedLab WMP checkpoint used as structural template.", + ) + parser.add_argument("--output", type=Path, required=True, help="Output converted checkpoint path.") + args = parser.parse_args() + + source = torch.load(args.source.expanduser().resolve(), map_location="cpu", weights_only=False) + reference = torch.load(args.reference.expanduser().resolve(), map_location="cpu", weights_only=False) + + required = ("model_state_dict", "world_model_dict", "depth_predictor") + missing = [key for key in required if key not in source] + if missing: + raise KeyError(f"Source is not an original WMP checkpoint, missing: {missing}") + + converted = dict(reference) + converted["actor_state_dict"] = _convert_actor(source["model_state_dict"], reference["actor_state_dict"]) + converted["world_model_state_dict"] = _copy_matching( + _convert_world_model(source["world_model_dict"]), + reference["world_model_state_dict"], + "world_model", + ) + converted["depth_predictor_state_dict"] = _copy_matching( + source["depth_predictor"], + reference["depth_predictor_state_dict"], + "depth_predictor", + ) + converted["iter"] = int(source.get("iter", reference.get("iter", 0))) + converted["infos"] = source.get("infos") + + output = args.output.expanduser().resolve() + output.parent.mkdir(parents=True, exist_ok=True) + torch.save(converted, output) + print(f"[INFO] converted checkpoint saved: {output}") + print("[INFO] converted modules: actor_state_dict, world_model_state_dict, depth_predictor_state_dict") + print("[WARN] critic/optimizer/AMP states are kept from reference template; inference uses converted actor + WMP modules.") + + +if __name__ == "__main__": + main() diff --git a/legged_lab/sensors/wmp_partial_tiled_camera.py b/legged_lab/sensors/wmp_partial_tiled_camera.py index 0c9eff9..f590c1f 100644 --- a/legged_lab/sensors/wmp_partial_tiled_camera.py +++ b/legged_lab/sensors/wmp_partial_tiled_camera.py @@ -27,48 +27,142 @@ def select_wmp_camera_env_ids( camera_num_envs: int | None = 1024, seed: int = 42, terrain_generator=None, + terrain_types: torch.Tensor | None = None, + terrain_cols_by_kind: dict[str, tuple[int, ...]] | None = None, force_tilt_crawl: bool = True, device: str | torch.device = "cpu", ) -> torch.Tensor: """选择 WMP 真实相机环境。 对齐原版公式： - depth_index = sample(non_tilt_crawl) U range(tilt_start_idx, crawl_end_idx) - 在 IsaacLab 中 terrain type 由 env index 按列分配，因此先由 terrain - proportions 还原每列的地形类型，再强制加入 tilt/crawl 列对应的 env。 + depth_index = sample(non_tilt_crawl) + all_tilt_crawl """ + _, depth_index, _ = select_wmp_depth_indices( + num_envs=num_envs, + camera_num_envs=camera_num_envs, + seed=seed, + terrain_generator=terrain_generator, + terrain_types=terrain_types, + terrain_cols_by_kind=terrain_cols_by_kind, + force_tilt_crawl=force_tilt_crawl, + device=device, + ) + return depth_index + + +def select_wmp_depth_indices( + num_envs: int, + camera_num_envs: int | None = 1024, + seed: int = 42, + terrain_generator=None, + terrain_types: torch.Tensor | None = None, + terrain_cols_by_kind: dict[str, tuple[int, ...]] | None = None, + force_tilt_crawl: bool = True, + device: str | torch.device = "cpu", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """生成原版 WMP 三元 depth index。""" device = torch.device(device) + inverse = torch.full((max(int(num_envs), 0),), -1, device=device, dtype=torch.long) if num_envs <= 0: - return torch.empty(0, device=device, dtype=torch.long) + empty = torch.empty(0, device=device, dtype=torch.long) + return empty, empty, inverse if camera_num_envs is None: camera_num_envs = 1024 camera_num_envs = min(int(camera_num_envs), int(num_envs)) - if camera_num_envs <= 0: - return torch.empty(0, device=device, dtype=torch.long) + camera_num_envs = max(camera_num_envs, 0) - forced = _tilt_crawl_env_ids(num_envs, terrain_generator, force_tilt_crawl, device) + forced = _tilt_crawl_env_ids( + num_envs, + terrain_generator, + force_tilt_crawl, + device, + terrain_types=terrain_types, + terrain_cols_by_kind=terrain_cols_by_kind, + ) + forced = _ordered_unique_valid(forced, num_envs, device) + if forced.numel() > camera_num_envs: + print( + "[WARN] WMP depth_index expands camera_num_envs to include all tilt/crawl envs: " + f"requested={camera_num_envs}, tilt_crawl={int(forced.numel())}" + ) + camera_num_envs = int(forced.numel()) remaining = camera_num_envs - int(forced.numel()) - if remaining <= 0: - return forced[:camera_num_envs].unique(sorted=True) - pool_mask = torch.ones(num_envs, device=device, dtype=torch.bool) - if forced.numel() > 0: - pool_mask[forced] = False - pool = pool_mask.nonzero(as_tuple=False).flatten() - if pool.numel() <= remaining: + pool = _depth_index_without_tilt_crawl_pool( + num_envs, + device, + terrain_types=terrain_types, + terrain_cols_by_kind=terrain_cols_by_kind, + forced=forced, + ) + if remaining <= 0: + sampled = torch.empty(0, device=device, dtype=torch.long) + elif pool.numel() <= remaining: sampled = pool else: generator = torch.Generator(device="cpu") generator.manual_seed(int(seed)) perm = torch.randperm(pool.numel(), generator=generator, device="cpu")[:remaining].to(device) sampled = pool[perm] - return torch.cat((forced, sampled), dim=0).unique(sorted=True) + depth_index_without_crawl_tilt = torch.sort(_ordered_unique_valid(sampled, num_envs, device)).values + forced = torch.sort(forced).values + depth_index = torch.cat((depth_index_without_crawl_tilt, forced), dim=0) + inverse[depth_index] = torch.arange(depth_index.numel(), device=device, dtype=torch.long) + return depth_index_without_crawl_tilt, depth_index, inverse + + +def _ordered_unique_valid(ids: torch.Tensor, num_envs: int, device: torch.device) -> torch.Tensor: + ids = torch.as_tensor(ids, device=device, dtype=torch.long).flatten() + if ids.numel() == 0: + return ids + keep = torch.zeros(num_envs, device=device, dtype=torch.bool) + out = [] + for env_id in ids.tolist(): + if 0 <= int(env_id) < num_envs and not bool(keep[int(env_id)].item()): + keep[int(env_id)] = True + out.append(int(env_id)) + return torch.tensor(out, device=device, dtype=torch.long) + + +def _depth_index_without_tilt_crawl_pool( + num_envs: int, + device: torch.device, + terrain_types: torch.Tensor | None, + terrain_cols_by_kind: dict[str, tuple[int, ...]] | None, + forced: torch.Tensor, +) -> torch.Tensor: + cols_by_kind = terrain_cols_by_kind or {} + forced_cols = list(cols_by_kind.get("tilt", ())) + list(cols_by_kind.get("crawl", ())) + if forced_cols and terrain_types is not None: + terrain_types = torch.as_tensor(terrain_types, device=device) + # 原版从 tilt_start_idx 之前采样，因此 rough_flat 等 crawl 后方地形不进入随机相机池。 + return (terrain_types < min(forced_cols)).nonzero(as_tuple=False).flatten() + pool_mask = torch.ones(num_envs, device=device, dtype=torch.bool) + if forced.numel() > 0: + pool_mask[forced] = False + return pool_mask.nonzero(as_tuple=False).flatten() -def _tilt_crawl_env_ids(num_envs: int, terrain_generator, force_tilt_crawl: bool, device: torch.device) -> torch.Tensor: - if not force_tilt_crawl or terrain_generator is None or len(terrain_generator.sub_terrains) <= 1: +def _tilt_crawl_env_ids( + num_envs: int, + terrain_generator, + force_tilt_crawl: bool, + device: torch.device, + terrain_types: torch.Tensor | None = None, + terrain_cols_by_kind: dict[str, tuple[int, ...]] | None = None, +) -> torch.Tensor: + if not force_tilt_crawl: return torch.empty(0, device=device, dtype=torch.long) + cols_by_kind = terrain_cols_by_kind or {} + forced_cols = list(cols_by_kind.get("tilt", ())) + list(cols_by_kind.get("crawl", ())) + if forced_cols and terrain_types is not None: + terrain_types = torch.as_tensor(terrain_types, device=device) + col_tensor = torch.tensor(forced_cols, device=device, dtype=terrain_types.dtype) + return torch.isin(terrain_types, col_tensor).nonzero(as_tuple=False).flatten() + + if terrain_generator is None or len(terrain_generator.sub_terrains) <= 1: + return torch.empty(0, device=device, dtype=torch.long) keys = list(terrain_generator.sub_terrains.keys()) forced_kind_ids = [idx for idx, key in enumerate(keys) if key.removeprefix("wmp_") in ("tilt", "crawl")] if not forced_kind_ids: @@ -96,7 +190,7 @@ class WMPPartialTiledCameraCfg(TiledCameraCfg): full_num_envs: int = 0 camera_env_ids: tuple[int, ...] = () - camera_model_name: str = "gemini2" + camera_model_name: str = "RGBD_camera" class WMPPartialTiledCamera(TiledCamera): @@ -113,7 +207,8 @@ def __init__(self, cfg: WMPPartialTiledCameraCfg): camera_ids = torch.as_tensor(cfg.camera_env_ids, dtype=torch.long) if camera_ids.numel() == 0: raise ValueError("WMPPartialTiledCamera requires at least one camera env id.") - self.camera_env_ids = camera_ids.unique(sorted=True) + ordered_num_envs = int(cfg.full_num_envs or int(camera_ids.max().item()) + 1) + self.camera_env_ids = _ordered_unique_valid(camera_ids, ordered_num_envs, torch.device("cpu")) full_num_envs = int(cfg.full_num_envs or (int(self.camera_env_ids.max().item()) + 1)) self.full_num_envs = full_num_envs self.camera_env_mask = torch.zeros(full_num_envs, dtype=torch.bool) @@ -138,11 +233,14 @@ def _initialize_impl(self): if match is not None: actual_env_ids.append(int(match.group(1))) if len(actual_env_ids) == self._view.count: - self.camera_env_ids = torch.tensor(actual_env_ids, dtype=torch.long) - self.camera_env_mask = torch.zeros(self.full_num_envs, dtype=torch.bool) - self.camera_env_mask[self.camera_env_ids] = True - self.camera_env_id_inverse = torch.full((self.full_num_envs,), -1, dtype=torch.long) - self.camera_env_id_inverse[self.camera_env_ids] = torch.arange(self.camera_env_ids.numel(), dtype=torch.long) + actual = torch.tensor(actual_env_ids, dtype=torch.long) + if actual.numel() != self.camera_env_ids.numel() or not torch.equal(actual, self.camera_env_ids): + requested_head = self.camera_env_ids[:16].tolist() + actual_head = actual[:16].tolist() + raise RuntimeError( + "WMPPartialTiledCamera image slot order must match depth_index exactly. " + f"requested_head={requested_head}, actual_head={actual_head}" + ) def _spawn_partial_cameras_from_cfg(self, cfg: WMPPartialTiledCameraCfg, spawn_cfg): stage = get_current_stage() diff --git a/legged_lab/utils/env_utils/scene.py b/legged_lab/utils/env_utils/scene.py index e20a557..d5315bd 100644 --- a/legged_lab/utils/env_utils/scene.py +++ b/legged_lab/utils/env_utils/scene.py @@ -89,7 +89,16 @@ def __init__(self, config: "BaseSceneCfg", physics_dt, step_dt): ) self.forward_height_scanner = RayCasterCfg( prim_path="{ENV_REGEX_NS}/Robot/" + config.height_scanner.prim_body_name, - offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), + # 原版 WMP forward height map 采样 x=[0,2], y=[-1.2,1.2]。 + # IsaacLab GridPattern 以局部原点为中心生成 [-1,1]x[-1.2,1.2]， + # 因此前移 1m 后得到 x' = x + 1 ∈ [0,2]。 + offset=RayCasterCfg.OffsetCfg( + pos=( + config.height_scanner.forward_offset[0], + config.height_scanner.forward_offset[1], + 20.0 + config.height_scanner.forward_offset[2], + ) + ), ray_alignment="yaw", pattern_cfg=patterns.GridPatternCfg( resolution=config.height_scanner.forward_resolution, diff --git a/legged_lab/world_models/wmp/controller.py b/legged_lab/world_models/wmp/controller.py index 7d92239..e6c9992 100644 --- a/legged_lab/world_models/wmp/controller.py +++ b/legged_lab/world_models/wmp/controller.py @@ -50,21 +50,29 @@ def __init__(self, env, config, world_model: WorldModel): lr=config.depth_predictor.lr, weight_decay=config.depth_predictor.weight_decay, ) - self.camera_env_ids = self._select_camera_env_ids() + self.depth_index = self._select_camera_env_ids() + self.camera_env_ids = self.depth_index + self.depth_index_without_crawl_tilt = self._select_depth_index_without_crawl_tilt() + self.depth_index_inverse = self._select_depth_index_inverse() self.camera_env_mask = torch.zeros(self.env.num_envs, device=self.device, dtype=torch.bool) if self.camera_env_ids.numel() > 0: self.camera_env_mask[self.camera_env_ids] = True + self.depth_predictor_excluded_env_ids = self._tilt_crawl_env_ids() max_episode_steps = int(getattr(self.env, "max_episode_length", 1000) // self.update_interval + 3) self.replay = WMPFixedEpisodeReplayBuffer( self.env.num_envs, max_episode_steps, - self.camera_env_ids, + self.depth_index, config.replay_device, + depth_index_without_crawl_tilt=self.depth_index_without_crawl_tilt, + depth_index_inverse=self.depth_index_inverse, + depth_predictor_excluded_env_ids=self.depth_predictor_excluded_env_ids, ) self.reset_state() def reset_state(self): self.latent = None + self.prior_latent = None self.is_first = torch.ones(self.env.num_envs, device=self.device) self.pending_obs: dict[str, torch.Tensor] | None = None self.step_counter = 0 @@ -94,7 +102,7 @@ def observe_before_policy(self) -> tuple[dict[str, torch.Tensor] | None, torch.T with torch.inference_mode(): embed = self.world_model.encoder(wm_obs) wm_action = self.action_history.reshape(self.env.num_envs, -1) - self.latent, _ = self.world_model.dynamics.obs_step( + self.latent, self.prior_latent = self.world_model.dynamics.obs_step( self.latent, wm_action, embed, @@ -104,6 +112,25 @@ def observe_before_policy(self) -> tuple[dict[str, torch.Tensor] | None, torch.T self.is_first.zero_() return wm_obs, self.feature + def predict_depth_image(self, use_prior: bool = True) -> torch.Tensor | None: + """从 RSSM latent 解码 depth image。 + + `use_prior=True` 对应 Dreamer/WMP 的一步预测： + prior_t = RSSM(s_{t-1}, a_{t-1}) + image_pred_t = Decoder(prior_t) + 用它和同一 WMP tick 的真实 image_t 对比，能看出世界模型是否真的在预测。 + """ + latent = self.prior_latent if use_prior else self.latent + if latent is None: + return None + with torch.inference_mode(): + feat = self.world_model.dynamics.get_feat(latent) + pred = self.world_model.decode(feat).get("image") + if pred is None: + return None + image = pred.mode() if callable(getattr(pred, "mode", None)) else pred.mean() + return image + def after_env_step( self, actions: torch.Tensor, @@ -151,8 +178,9 @@ def after_env_step( self.step_counter = next_step_counter def train_if_ready(self, iteration: int, total_env_steps: int) -> dict[str, float]: + del total_env_steps metrics = {} - if total_env_steps < self.config.train_start_steps: + if getattr(self.replay, "num_steps", 0) <= self.config.train_start_steps: return metrics train_interval = max(1, int(getattr(self.config, "train_interval", 1))) if iteration % train_interval != 0: @@ -174,6 +202,11 @@ def replay_stats(self) -> dict[str, float]: stats = { "wm_replay_episodes": float(len(self.replay)), "wm_replay_steps": float(getattr(self.replay, "num_steps", 0)), + "wm_depth_predictor_excluded_envs": float(self.depth_predictor_excluded_env_ids.numel()), + "wm_depth_index_count": float(self.depth_index.numel()), + "wm_depth_index_without_crawl_tilt_count": float(self.depth_index_without_crawl_tilt.numel()), + "wm_tilt_crawl_depth_count": float(self.depth_index.numel() - self.depth_index_without_crawl_tilt.numel()), + "wm_real_depth_ratio": float(self.depth_index.numel()) / max(float(self.env.num_envs), 1.0), } if hasattr(self.replay, "max_episode_steps"): stats["wm_replay_max_episode_steps"] = float(self.replay.max_episode_steps) @@ -192,6 +225,10 @@ def load_state_dict(self, state: dict, strict: bool = True, load_optimizer: bool self.depth_predictor_opt.load_state_dict(state["depth_predictor_optimizer_state_dict"]) def _read_wm_obs(self) -> dict[str, torch.Tensor]: + # 调用点在 `after_env_step()` 中，且只在 env.step 完成后的 WMP tick 读取。 + # A1 WMP 配置里 camera.update_period 与 sim.render_interval 都等于 + # step_dt * update_interval，因此这里的 prop / forward_height_map / depth + # 对应同一个完成后的机器人状态。 prop = self.env.get_wmp_proprioception().to(self.device) forward_height_map = self.env.get_wmp_forward_height_map().to(self.device) image, has_real_depth = self._current_depth_image(prop, forward_height_map) @@ -221,18 +258,26 @@ def _current_depth_image(self, prop: torch.Tensor, forward_height_map: torch.Ten 这对应原版 WMP 的公式： image_all = DepthPredictor(forward_height_map, prop) - image_all[camera_env_ids] = real_depth[camera_env_ids] + image_all[depth_index] = real_depth """ with torch.inference_mode(): image = self.depth_predictor(forward_height_map, prop) has_real_depth = torch.zeros(self.env.num_envs, device=self.device) + if hasattr(self.env, "get_wmp_depth_observations") and getattr(self.env, "wmp_depth_buffer", None) is not None: + real_image = self.env.get_wmp_depth_observations().to(self.device) + real_env_ids = self._real_depth_env_ids(real_image.shape[0]) + if real_env_ids.numel() > 0: + image = image.clone() + image[real_env_ids] = real_image[: real_env_ids.numel()] + has_real_depth[real_env_ids] = 1.0 + return image, has_real_depth if hasattr(self.env, "get_depth_observations"): try: depth = self.env.get_depth_observations().to(self.device) real_image = depth_to_wmp_image( depth, - near=self.env.cfg.scene.gemini2_camera.depth_near, - far=self.env.cfg.scene.gemini2_camera.depth_far, + near=self.env.cfg.scene.rgbd_camera.depth_near, + far=self.env.cfg.scene.rgbd_camera.depth_far, ) real_env_ids = self._real_depth_env_ids(real_image.shape[0]) if real_env_ids.numel() > 0: @@ -247,24 +292,34 @@ def _current_depth_image(self, prop: torch.Tensor, forward_height_map: torch.Ten return image, has_real_depth def _real_depth_env_ids(self, depth_batch_size: int) -> torch.Tensor: + if hasattr(self.env, "get_depth_index"): + env_ids = self.env.get_depth_index().to(self.device).long() + if depth_batch_size == self.env.num_envs: + return env_ids[(env_ids >= 0) & (env_ids < self.env.num_envs)] + return env_ids[: min(depth_batch_size, env_ids.numel())] if hasattr(self.env, "get_depth_camera_env_ids"): env_ids = self.env.get_depth_camera_env_ids().to(self.device).long() if depth_batch_size == self.env.num_envs: return env_ids[(env_ids >= 0) & (env_ids < self.env.num_envs)] - return env_ids[:depth_batch_size] + return env_ids[: min(depth_batch_size, env_ids.numel())] if depth_batch_size == self.env.num_envs: return self.camera_env_ids - return self.camera_env_ids[:depth_batch_size] + return self.camera_env_ids[: min(depth_batch_size, self.camera_env_ids.numel())] def _select_camera_env_ids(self) -> torch.Tensor: if self.config.camera_env_ids is not None: ids = torch.as_tensor(self.config.camera_env_ids, device=self.device, dtype=torch.long) - return ids[(ids >= 0) & (ids < self.env.num_envs)].unique(sorted=True) + return self._ordered_unique_valid(ids) + + if hasattr(self.env, "get_depth_index"): + ids = self.env.get_depth_index().to(self.device).long() + if ids.numel() > 0 and ids.numel() <= self.env.num_envs: + return self._ordered_unique_valid(ids) if hasattr(self.env, "get_depth_camera_env_ids"): ids = self.env.get_depth_camera_env_ids().to(self.device).long() if ids.numel() > 0 and ids.numel() < self.env.num_envs: - return ids.unique(sorted=True) + return self._ordered_unique_valid(ids) if self.config.camera_sample_all_envs: return torch.arange(self.env.num_envs, device=self.device, dtype=torch.long) @@ -273,13 +328,18 @@ def _select_camera_env_ids(self) -> torch.Tensor: if camera_num_envs is None: camera_num_envs = 1024 camera_num_envs = min(int(camera_num_envs), self.env.num_envs) - if camera_num_envs <= 0: - return torch.empty(0, device=self.device, dtype=torch.long) + camera_num_envs = max(camera_num_envs, 0) forced = self._forced_camera_env_ids() + if forced.numel() > camera_num_envs: + print( + "[WARN] WMP depth_index expands camera_num_envs to include all tilt/crawl envs: " + f"requested={camera_num_envs}, tilt_crawl={int(forced.numel())}" + ) + camera_num_envs = int(forced.numel()) remaining = camera_num_envs - int(forced.numel()) if remaining <= 0: - return forced[:camera_num_envs].unique(sorted=True) + return self._ordered_unique_valid(forced) pool_mask = torch.ones(self.env.num_envs, device=self.device, dtype=torch.bool) if forced.numel() > 0: @@ -292,33 +352,79 @@ def _select_camera_env_ids(self) -> torch.Tensor: generator.manual_seed(int(self.config.camera_env_seed)) perm = torch.randperm(pool.numel(), generator=generator, device="cpu")[:remaining].to(self.device) sampled = pool[perm] - return torch.cat((forced, sampled), dim=0).unique(sorted=True) + return self._ordered_unique_valid(torch.cat((sampled, forced), dim=0)) + + def _select_depth_index_without_crawl_tilt(self) -> torch.Tensor: + if hasattr(self.env, "get_depth_index_without_crawl_tilt"): + ids = self.env.get_depth_index_without_crawl_tilt().to(self.device).long() + if ids.numel() > 0: + return self._ordered_unique_valid(ids) + if self.depth_index.numel() == 0: + return torch.empty(0, device=self.device, dtype=torch.long) + tilt_crawl = self._tilt_crawl_env_ids() + if tilt_crawl.numel() == 0: + return self.depth_index + keep_mask = ~torch.isin(self.depth_index, tilt_crawl) + return self.depth_index[keep_mask] + + def _select_depth_index_inverse(self) -> torch.Tensor: + if hasattr(self.env, "get_depth_index_inverse"): + inverse = self.env.get_depth_index_inverse().to(self.device).long() + if inverse.numel() == self.env.num_envs: + return inverse + inverse = torch.full((self.env.num_envs,), -1, device=self.device, dtype=torch.long) + if self.depth_index.numel() > 0: + inverse[self.depth_index] = torch.arange(self.depth_index.numel(), device=self.device, dtype=torch.long) + return inverse + + def _ordered_unique_valid(self, ids: torch.Tensor) -> torch.Tensor: + ids = torch.as_tensor(ids, device=self.device, dtype=torch.long).flatten() + keep = torch.zeros(self.env.num_envs, device=self.device, dtype=torch.bool) + out = [] + for env_id in ids.tolist(): + if 0 <= int(env_id) < self.env.num_envs and not bool(keep[int(env_id)].item()): + keep[int(env_id)] = True + out.append(int(env_id)) + return torch.tensor(out, device=self.device, dtype=torch.long) def _forced_camera_env_ids(self) -> torch.Tensor: if not self.config.camera_force_tilt_crawl: return torch.empty(0, device=self.device, dtype=torch.long) - terrain_types = getattr(self.env, "terrain_types", None) - terrain_generator = getattr(self.env.cfg.scene, "terrain_generator", None) - if terrain_types is None or terrain_generator is None or len(terrain_generator.sub_terrains) <= 1: - return torch.empty(0, device=self.device, dtype=torch.long) + return self._tilt_crawl_env_ids() - keys = list(terrain_generator.sub_terrains.keys()) - tilt_cols = [idx for idx, key in enumerate(keys) if key.removeprefix("wmp_") in ("tilt", "crawl")] - if not tilt_cols: + def _tilt_crawl_env_ids(self) -> torch.Tensor: + terrain_types = getattr(self.env, "terrain_types", None) + if terrain_types is None: return torch.empty(0, device=self.device, dtype=torch.long) - proportions = torch.tensor( - [float(terrain_generator.sub_terrains[key].proportion) for key in keys], - device=self.device, - ) - proportions = proportions / proportions.sum().clamp_min(1.0e-6) - cumulative = torch.cumsum(proportions, dim=0) - num_cols = int(terrain_generator.num_cols) - col_kind = [] - for col in range(num_cols): - kind = int(torch.searchsorted(cumulative, torch.tensor(col / num_cols + 0.001, device=self.device), right=False).item()) - col_kind.append(kind) - tilt_crawl_cols = [col for col, kind in enumerate(col_kind) if kind in tilt_cols] + cols_by_kind = getattr(self.env, "wmp_terrain_cols_by_kind", {}) or {} + tilt_crawl_cols = list(cols_by_kind.get("tilt", ())) + list(cols_by_kind.get("crawl", ())) + if not tilt_crawl_cols: + terrain_generator = getattr(self.env.cfg.scene, "terrain_generator", None) + if terrain_generator is None or len(terrain_generator.sub_terrains) <= 1: + return torch.empty(0, device=self.device, dtype=torch.long) + keys = list(terrain_generator.sub_terrains.keys()) + tilt_kind_ids = [idx for idx, key in enumerate(keys) if key.removeprefix("wmp_") in ("tilt", "crawl")] + if not tilt_kind_ids: + return torch.empty(0, device=self.device, dtype=torch.long) + proportions = torch.tensor( + [float(terrain_generator.sub_terrains[key].proportion) for key in keys], + device=self.device, + ) + proportions = proportions / proportions.sum().clamp_min(1.0e-6) + cumulative = torch.cumsum(proportions, dim=0) + num_cols = int(terrain_generator.num_cols) + col_kind = [] + for col in range(num_cols): + kind = int( + torch.searchsorted( + cumulative, + torch.tensor(col / num_cols + 0.001, device=self.device), + right=False, + ).item() + ) + col_kind.append(kind) + tilt_crawl_cols = [col for col, kind in enumerate(col_kind) if kind in tilt_kind_ids] if not tilt_crawl_cols: return torch.empty(0, device=self.device, dtype=torch.long) @@ -356,19 +462,33 @@ def _train_depth_predictor(self) -> float | None: return total_loss / used_iters def _train_world_model(self) -> dict[str, float]: - metrics = {} + metrics_sum = {} + used_steps = 0 for _ in range(self.config.train_steps_per_iter): batch = self.replay.sample(self.config.batch_size, self.config.batch_length) batch = self._materialize_sampled_images(batch) - batch = {k: v for k, v in batch.items() if k not in ("forward_height_map", "has_real_depth")} + batch = { + k: v + for k, v in batch.items() + if k not in ("forward_height_map", "has_real_depth", "depth_index_slot") + } _, _, metrics = self.world_model._train(batch) - return metrics + for key, value in metrics.items(): + metrics_sum[key] = metrics_sum.get(key, 0.0) + float(value) + used_steps += 1 + if used_steps == 0: + return {} + return {key: value / used_steps for key, value in metrics_sum.items()} def _materialize_sampled_images(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - has_real_depth = batch.get("has_real_depth") - if has_real_depth is None: - return batch - missing_mask = has_real_depth[..., 0] <= 0.5 + depth_index_slot = batch.get("depth_index_slot") + if depth_index_slot is None: + has_real_depth = batch.get("has_real_depth") + if has_real_depth is None: + return batch + missing_mask = has_real_depth[..., 0] <= 0.5 + else: + missing_mask = depth_index_slot[..., 0] < 0 if not torch.any(missing_mask): return batch From d3f3633794a2b62ca9427f1044979c12b5fe3df3 Mon Sep 17 00:00:00 2001 From: m0_51228868 Date: Tue, 9 Jun 2026 09:28:10 +0800 Subject: [PATCH 2/6] align WMP training semantics --- WMP_NETWORK_DESIGN_DIFF_REPORT.txt | 231 ++++++++++++++++++++++++++- legged_lab/algorithms/wmp_amp_ppo.py | 42 +++-- legged_lab/envs/a1/a1_config.py | 10 +- legged_lab/envs/base/base_config.py | 7 + legged_lab/envs/base/base_env.py | 34 +++- legged_lab/runners/wmp_amp_runner.py | 28 ++++ 6 files changed, 333 insertions(+), 19 deletions(-) diff --git a/WMP_NETWORK_DESIGN_DIFF_REPORT.txt b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt index 5b4415a..fcdaefb 100644 --- a/WMP_NETWORK_DESIGN_DIFF_REPORT.txt +++ b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt @@ -551,7 +551,7 @@ reward 公式: entropy_coef=0.01 num_learning_epochs=5 num_mini_batches=4 - normalize_advantage_per_mini_batch=True + normalize_advantage_per_mini_batch=False vel_predict_coef=1.0 num_steps_per_env=24 @@ -648,7 +648,6 @@ critic 维度已经对齐，下一步不是改维度，而是验收字段语义: 建议做: - actor obs_normalization 开关实验。 - critic obs_normalization 开关实验。 - - normalize_advantage_per_mini_batch 开关实验。 - desired_kl/schedule 与原版配置对照实验。 @@ -676,3 +675,231 @@ critic 维度已经对齐，下一步不是改维度，而是验收字段语义: IsaacLab 与 IsaacGym 的仿真和传感器差异； AMP retarget 后 expert 数据分布； rsl_rl v5 PPO/storage/normalization 更新接口差异。 + + +12. 2026-06-09 复查新增的剩余差异 +================================= + +本节是在前述 shape/语义修复之后，继续对比原版 WMP 当前仍存在的差异。 +结论先行: + 网络主干和关键输入维度已经不是主要差异； + 剩余差异主要集中在训练循环、PPO/AMP 优化接口、terrain/camera 实现和 replay 采样细节； + A1 WMP 专用 reset/termination 语义已在本轮对齐，接触传感器实现仍保留 IsaacLab 框架差异。 + + +12.1 PPO / AMP optimizer 仍不同 +------------------------------- +原版 AMPPPO: + - actor_critic 和 discriminator 由同一个 Adam optimizer 管理。 + - optimizer param groups: + actor_critic + amp_trunk(weight_decay=10e-4) + amp_head(weight_decay=10e-2) + - 每次 update 后只对 actor_critic.parameters() 做 grad clip。 + - 若 min_normalized_std 存在，会对 actor_critic.std 做下限 clamp。 + +当前 WMPAMPPPO: + - 继承 rsl_rl v5 PPO，actor 和 critic 是两个模型对象。 + - discriminator 通过 add_param_group 挂到 rsl_rl v5 optimizer。 + - update 时分别 clip actor.parameters() 与 critic.parameters()。 + - 已复刻原版 min_normalized_std clamp: + A1 配置为 [0.05, 0.02, 0.05] * 4； + WMPAMPRunner 使用 IsaacLab joint soft limits 计算 + min_std = min_normalized_std * (upper_limit - lower_limit)； + WMPAMPPPO 在 optimizer.step() 后 clamp GaussianDistribution 的 std 参数。 + +影响: + 即使 loss 公式接近，优化器参数组织、grad clip 范围和 rsl_rl v5 storage/update 细节仍会改变 PPO 轨迹。 + + +12.2 Advantage / return / timeout 已按原版语义对齐 +-------------------------------------------------- +原版: + - AMPPPO.process_env_step() 会处理 infos["time_outs"] bootstrap: + reward += gamma * value * time_out + - storage.compute_returns() 在整段 rollout 上计算 GAE。 + - advantages 在 compute_returns() 后按全 rollout 做一次标准化。 + +当前: + - BaseEnv.step() 每步写入 extras["time_outs"] = self.time_out_buf。 + - WMPAMPPPO.process_env_step() 调用 rsl_rl v5 PPO.process_env_step()，父类已经执行 timeout bootstrap。 + - WMPAMPPPO 不再覆盖 compute_returns()，直接使用 rsl_rl v5 PPO.compute_returns()。 + - A1 WMP 配置已设 normalize_advantage_per_mini_batch=False，因此 advantages 按全 rollout 标准化。 + +影响: + timeout bootstrap、GAE 公式和 advantage 标准化位置已经与原版语义对齐。 + 剩余差异仅来自 rsl_rl v5 storage/TensorDict 的实现细节，而不是算法公式差异。 + + +12.3 Reset / termination 已按原版 WMP A1 语义对齐 +------------------------------------------------- +原版 LeggedRobot.check_termination(): + - base 触地终止: + reset_contact = any(norm(contact_forces[:, termination_contact_indices, :]) > 1) + - time_out 终止: + time_out_buf = episode_length_buf > max_episode_length + - vel_violate 终止: + vel_error = base_lin_vel[:, 0] - commands[:, 0] + vel_violate = + ((vel_error > 1.5) & (commands[:, 0] < 0)) | + ((vel_error < -1.5) & (commands[:, 0] > 0)) + vel_violate &= terrain_levels > 3 + - fall 终止: + fall = (root_states[:, 9] < -3) | (projected_gravity[:, 2] > 0) + - reset_buf = reset_contact | time_out_buf | vel_violate | fall。 + +当前 A1 WMP BaseEnv.check_reset(): + - base 触地终止: + terminate_contacts_body_names = [".*trunk.*"]，阈值仍为 > 1。 + 当前使用 IsaacLab ContactSensor 的 net_forces_w_history 历史窗口取 max； + 原版使用 IsaacGym 当前步 contact_forces。这是保留的框架层实现差异。 + - time_out 终止: + A1 WMP 设置 wmp_time_out_strictly_greater=True， + 因此使用 episode_length_buf > max_episode_length，与原版一致。 + 其他任务默认仍保持原来的 >= 语义，避免扩大行为变更。 + - vel_violate 终止: + terminate_on_wmp_velocity_violation=True，仅 A1 WMP 开启。 + vel_error = robot.data.root_lin_vel_b[:, 0] - command_generator.command[:, 0] + 阈值 wmp_velocity_violation_threshold = 1.5。 + wmp_velocity_violation_min_terrain_level = 4，因此条件等价 terrain_levels > 3。 + 若当前环境没有 terrain_levels，则 vel_violate 保持 False，避免平地任务误触发。 + - fall 终止: + terminate_on_wmp_fall=True，仅 A1 WMP 开启。 + robot.data.root_lin_vel_w[:, 2] < -3.0 对齐原版 root_states[:, 9] < -3。 + robot.data.projected_gravity_b[:, 2] > 0.0 对齐原版 projected_gravity[:, 2] > 0。 + - reset cause 日志: + extras["log"]["wmp_vel_violate"] 记录本步速度违例比例。 + extras["log"]["wmp_fall"] 记录本步 fall 比例。 + +影响: + A1 WMP 的 episode 长度分布、AMP terminal state、timeout bootstrap 标志和 terrain curriculum reset 入口， + 已按原版 WMP A1 的主要 reset 公式对齐。 + +残余差异: + - 接触终止的数据源仍是 IsaacLab ContactSensor 的历史窗口，而不是 IsaacGym 当前步 contact_forces。 + 这可能让 trunk 短暂接触的 reset 更保守/更敏感，后续可通过 reset cause 统计决定是否再增加 + WMP 专用 current-force 路径。 + - reset 后随机化、scene.reset 和 command reset 的内部执行顺序仍遵循当前 IsaacLab BaseEnv， + 不是 legged_gym reset_idx 的逐行迁移。 + + +12.4 Camera pose / pitch 已人工审核通过 +-------------------------------------- +原版 A1 WMP depth: + - camera position = [0.27, 0, 0.03] + - y_angle = [-5, 5]，注释为 positive pitch down + - horizontal_fov = 58 + - near/far = 0/2 + - buffer_len = 2 + +当前 A1 WMP: + - spawn_offset_pos = (0.3, 0.0, 0.03) + - random_pitch_deg = (5.0, 15.0) + - horizontal_fov_deg = 58.0 + - depth_near/depth_far = 0/2 + - WMP depth buffer len = 2 + +已对齐: + - FOV、near/far、buffer_len、depth_buffer[:, -2] 语义已经接近。 + +结论: + - 用户已人工审核当前 camera 外参和 pitch 设置，确认无需继续修订。 + - 因此 camera pose / pitch 不再列为当前 WMP diff 的优先待办。 + - 后续只需在真实 depth 分布异常时，再回到相机图像统计和可视化检查。 + + +12.5 World model replay 采样仍不是逐行等价 +------------------------------------------ +原版 train_world_model(): + - p = wm_dataset_size / sum(wm_dataset_size) 按 env episode 长度加权采样 env。 + - batch_length = min(sampled_env_dataset_size.min(), config.batch_length)。 + - 对每个 env 采样 end_idx，再取 [end_idx-batch_length, end_idx)。 + - 非真实相机 env 的 image 在 sample 时即时由 DepthPredictor(forward_height_map, prop) 生成。 + - is_first 每个 batch 第一帧置 1。 + +当前 WMPFixedEpisodeReplayBuffer: + - 也按 dataset_size 加权采样 env，并把 batch 第一帧 is_first 置 1。 + - 但 replay 存储、episode finish、image 缺失标记、depth_index_slot 和 sample_real_depth 的实现不是原版逐行迁移。 + - 当前 train_start_steps 使用 replay.num_steps；原版使用 sum_wm_dataset_size。 + +影响: + 当前方向是语义接近，但仍应把 world model loss 曲线差异优先归因到 replay/tick/sample 细节，而不是 RSSM 网络结构。 + + +12.6 DepthPredictor 监督采样仍有细节差异 +---------------------------------------- +原版 train_depth_predictor(): + - batch_idx 从 env.depth_index_without_crawl_tilt 采样。 + - time_index 对每个 env 从 [0, wm_dataset_size[idx]] 采样。 + - image 使用 env.depth_index_inverse[batch_idx] 映射到真实相机 image buffer。 + +当前: + - sample_real_depth() 从 camera_env_ids 中 dataset_size > 0 的 env 采样。 + - 是否严格排除 tilt/crawl 取决于 camera_env_ids、depth_index_without_crawl_tilt 和 has_real_depth 的联动。 + +风险: + 如果 tilt/crawl env 进入 DepthPredictor 监督采样，训练分布会比原版更难； + 如果 time_index 边界和原版不同，也会改变 early episode 的监督样本。 + + +12.7 Domain randomization 数值来源接近，但执行时机不同 +------------------------------------------------------ +原版: + - friction/restitution/base mass/link mass/COM/gains 在 IsaacGym 创建 env 或 reset callback 中随机。 + - motor_strength 在每个 sim substep 对 torque 乘随机系数。 + - action_latency 每次 step 随机整数延迟。 + +当前: + - friction/restitution/base mass/COM/gains 使用 IsaacLab EventTerm 或记录型 EventTerm。 + - motor_strength 在 reset 时采样，并作用于 action target。 + - action_delay 使用 IsaacLab DelayBuffer / sim_step 语义。 + +影响: + critic privileged buffer 已尽量记录对应随机量，但物理执行时机、作用位置和随机重采样频率仍不等价。 + + +12.8 Reward / command 基本对齐，但仍需跑统计 +------------------------------------------- +已对齐点: + - command x/y/yaw/heading 范围与原版 A1 WMP 接近。 + - reward curriculum: feet_edge 从 0.1 到 1.0，区间 4000 到 10000。 + - tracking_lin_vel、tracking_ang_vel、feet_edge、dof_error 等主要权重在 A1WMPAMPRewardCfg 中接近原版。 + +仍需关注: + - IsaacLab reward 函数实现不是 legged_gym 原函数逐行拷贝。 + - feet_edge 使用当前 terrain edge query；原版使用 legged_gym terrain 数据结构。 + - command generator、heading control 和 terrain curriculum 的时序来自 IsaacLab 当前实现。 + + +12.9 保存/恢复 checkpoint 格式不同 +---------------------------------- +原版 WMP runner save(): + - model_state_dict + - optimizer_state_dict + - world_model_dict + - wm_optimizer_state_dict + - depth_predictor + - iter + +当前 WMPAMPRunner save(): + - rsl_rl v5 actor/critic/distribution/optimizer 状态； + - amp_discriminator_state_dict、amp_normalizer_state_dict； + - world_model_state_dict、world_model_optimizer_state_dict； + - depth_predictor_state_dict、depth_predictor_optimizer_state_dict。 + +影响: + 当前 checkpoint 更完整，但和原版 checkpoint key/optimizer/state 结构不兼容。 + 旧 pt 需要转换脚本；从零训练不受此限制。 + + +12.10 复查后的优先级更新 +------------------------ +如果下一步要继续缩小与原版训练效果的差距，建议优先级为: + + 1. 对比一批 WMP replay sample: + prop/action/reward/is_first/image/forward_height_map 的 shape、时间索引、真实 depth 覆盖位置。 + 2. 对比 DepthPredictor 监督采样 env ids，确保默认只用 depth_index_without_crawl_tilt。 + 3. 做 PPO 接口 ablation: + obs normalizer、rsl_rl v5 storage/update 细节。 + 4. 记录 reset cause 分布、episode length 分布、AMP policy/expert score 分布。 + 5. 如果 trunk contact reset 占比异常，再评估是否为 A1 WMP 增加 current-force 接触终止路径。 diff --git a/legged_lab/algorithms/wmp_amp_ppo.py b/legged_lab/algorithms/wmp_amp_ppo.py index d1565f8..66592ed 100644 --- a/legged_lab/algorithms/wmp_amp_ppo.py +++ b/legged_lab/algorithms/wmp_amp_ppo.py @@ -25,6 +25,29 @@ def __init__( self.vel_predict_coef = float(vel_predict_coef) self.vel_target_start = int(vel_target_start) self.vel_target_dim = int(vel_target_dim) + self.min_action_std: torch.Tensor | None = None + + def set_min_action_std(self, min_action_std: torch.Tensor | None) -> None: + if min_action_std is None: + self.min_action_std = None + return + self.min_action_std = min_action_std.detach().to(self.device).float() + + def _clamp_action_std(self) -> None: + if self.min_action_std is None: + return + distribution = getattr(self.actor, "distribution", None) + if distribution is None: + return + # 原版 AMPPPO 在每次 optimizer.step() 后执行: + # std = clamp(std, min_normalized_std * (upper_limit - lower_limit)) + # rsl_rl v5 的 GaussianDistribution 支持 scalar std 或 log std 两种参数化。 + if hasattr(distribution, "std_param"): + min_std = self.min_action_std.to(distribution.std_param.device) + distribution.std_param.data.clamp_(min=min_std) + elif hasattr(distribution, "log_std_param"): + min_log_std = torch.log(torch.clamp(self.min_action_std, min=1.0e-8)).to(distribution.log_std_param.device) + distribution.log_std_param.data.clamp_(min=min_log_std) def attach_amp( self, @@ -58,23 +81,11 @@ def process_env_step(self, obs, rewards, dones, extras, next_amp_obs=None, task_ self._last_task_reward = task_rewards.detach().to(self.device) super().process_env_step(obs, rewards, dones, extras) - def compute_returns(self, obs): - st = self.storage - last_values = self.critic(obs).detach() - advantage = 0 - for step in reversed(range(st.num_transitions_per_env)): - next_values = last_values if step == st.num_transitions_per_env - 1 else st.values[step + 1] - next_is_not_terminal = 1.0 - st.dones[step].float() - delta = st.rewards[step] + next_is_not_terminal * self.gamma * next_values - st.values[step] - advantage = delta + next_is_not_terminal * self.gamma * self.lam * advantage - st.returns[step] = advantage + st.values[step] - st.advantages = st.returns - st.values - if not self.normalize_advantage_per_mini_batch: - st.advantages = (st.advantages - st.advantages.mean()) / (st.advantages.std(unbiased=False) + 1e-8) - def update(self): if self.amp_storage.num_samples == 0: - return super().update() + loss_dict = super().update() + self._clamp_action_std() + return loss_dict mean_value_loss = 0.0 mean_surrogate_loss = 0.0 @@ -180,6 +191,7 @@ def update(self): if self.is_multi_gpu: self.reduce_parameters() self.optimizer.step() + self._clamp_action_std() self.amp_normalizer.update(policy_state) self.amp_normalizer.update(expert_state) diff --git a/legged_lab/envs/a1/a1_config.py b/legged_lab/envs/a1/a1_config.py index d33e62c..fcdb528 100644 --- a/legged_lab/envs/a1/a1_config.py +++ b/legged_lab/envs/a1/a1_config.py @@ -230,6 +230,13 @@ def __post_init__(self): self.robot.actor_obs_history_length = 5 self.robot.critic_obs_history_length = 1 self.robot.wmp_privileged_contact_body_names = [".*thigh.*", ".*calf.*"] + self.robot.wmp_time_out_strictly_greater = True + self.robot.terminate_on_wmp_velocity_violation = True + self.robot.wmp_velocity_violation_threshold = 1.5 + self.robot.wmp_velocity_violation_min_terrain_level = 4 + self.robot.terminate_on_wmp_fall = True + self.robot.wmp_fall_z_velocity_threshold = -3.0 + self.robot.wmp_fall_projected_gravity_z_threshold = 0.0 self.normalization.obs_scales.ang_vel = 0.25 self.normalization.obs_scales.joint_pos = 1.0 self.normalization.obs_scales.joint_vel = 0.05 @@ -365,6 +372,7 @@ class A1WMPAMPTerrainAgentCfg(A1AMPFlatAgentCfg): "replay_buffer_size": 1000000, "preload_normalizer": True, "grad_penalty_coef": 1.0, + "min_normalized_std": [0.05, 0.02, 0.05] * 4, } def __post_init__(self): @@ -379,7 +387,7 @@ def __post_init__(self): self.algorithm.entropy_coef = 0.01 self.algorithm.num_learning_epochs = 5 self.algorithm.num_mini_batches = 4 - self.algorithm.normalize_advantage_per_mini_batch = True + self.algorithm.normalize_advantage_per_mini_batch = False self.algorithm.vel_predict_coef = 1.0 self.algorithm.vel_target_start = 50 self.algorithm.vel_target_dim = 3 diff --git a/legged_lab/envs/base/base_config.py b/legged_lab/envs/base/base_config.py index a9fcae7..cd19f1e 100644 --- a/legged_lab/envs/base/base_config.py +++ b/legged_lab/envs/base/base_config.py @@ -130,6 +130,13 @@ class RobotCfg: wmp_privileged_contact_body_names: list = [] terminate_on_flight: bool = False terminate_on_flight_threshold: float = 1.0 + wmp_time_out_strictly_greater: bool = False + terminate_on_wmp_velocity_violation: bool = False + wmp_velocity_violation_threshold: float = 1.5 + wmp_velocity_violation_min_terrain_level: int = 4 + terminate_on_wmp_fall: bool = False + wmp_fall_z_velocity_threshold: float = -3.0 + wmp_fall_projected_gravity_z_threshold: float = 0.0 @configclass diff --git a/legged_lab/envs/base/base_env.py b/legged_lab/envs/base/base_env.py index 0074f29..7adad11 100644 --- a/legged_lab/envs/base/base_env.py +++ b/legged_lab/envs/base/base_env.py @@ -43,6 +43,8 @@ def __init__(self, cfg: BaseEnvCfg, headless): self.depth_index = torch.empty(0, device=self.device, dtype=torch.long) self.depth_index_inverse = torch.empty(0, device=self.device, dtype=torch.long) self.wmp_depth_buffer = None + self.wmp_vel_violate_buf = None + self.wmp_fall_buf = None self.physics_dt = self.cfg.sim.dt self.step_dt = self.cfg.sim.decimation * self.cfg.sim.dt render_interval = self.cfg.sim.render_interval @@ -504,12 +506,17 @@ def step(self, actions: torch.Tensor): env_ids = self.reset_buf.nonzero(as_tuple=False).flatten() terminal_amp_states = self.get_amp_observations()[env_ids] self.reset(env_ids) + if self.cfg.robot.terminate_on_wmp_velocity_violation and self.wmp_vel_violate_buf is not None: + self.extras.setdefault("log", {})["wmp_vel_violate"] = self.wmp_vel_violate_buf.float().mean() + if self.cfg.robot.terminate_on_wmp_fall and self.wmp_fall_buf is not None: + self.extras.setdefault("log", {})["wmp_fall"] = self.wmp_fall_buf.float().mean() actor_obs, critic_obs = self.compute_observations() obs = TensorDict({"policy": actor_obs, "critic": critic_obs}, batch_size=[self.num_envs]) self.extras["observations"] = {"critic": critic_obs} self.extras["reset_env_ids"] = env_ids self.extras["terminal_amp_states"] = terminal_amp_states + self.extras["time_outs"] = self.time_out_buf return obs, reward_buf, self.reset_buf, self.extras @@ -679,7 +686,32 @@ def check_reset(self): > self.cfg.robot.terminate_on_flight_threshold ) reset_buf |= torch.sum(feet_contact, dim=-1) < 0.5 - time_out_buf = self.episode_length_buf >= self.max_episode_length + vel_violate = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + if self.cfg.robot.terminate_on_wmp_velocity_violation: + vel_error = self.robot.data.root_lin_vel_b[:, 0] - self.command_generator.command[:, 0] + threshold = float(self.cfg.robot.wmp_velocity_violation_threshold) + vel_violate = ((vel_error > threshold) & (self.command_generator.command[:, 0] < 0.0)) | ( + (vel_error < -threshold) & (self.command_generator.command[:, 0] > 0.0) + ) + terrain_levels = getattr(self, "terrain_levels", None) + if terrain_levels is not None: + vel_violate &= terrain_levels > int(self.cfg.robot.wmp_velocity_violation_min_terrain_level) - 1 + else: + vel_violate &= False + reset_buf |= vel_violate + fall = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + if self.cfg.robot.terminate_on_wmp_fall: + # 原版 WMP: fall = (root_states[:, 9] < -3) | (projected_gravity[:, 2] > 0)。 + fall = (self.robot.data.root_lin_vel_w[:, 2] < float(self.cfg.robot.wmp_fall_z_velocity_threshold)) | ( + self.robot.data.projected_gravity_b[:, 2] > float(self.cfg.robot.wmp_fall_projected_gravity_z_threshold) + ) + reset_buf |= fall + self.wmp_vel_violate_buf = vel_violate + self.wmp_fall_buf = fall + if self.cfg.robot.wmp_time_out_strictly_greater: + time_out_buf = self.episode_length_buf > self.max_episode_length + else: + time_out_buf = self.episode_length_buf >= self.max_episode_length reset_buf |= time_out_buf return reset_buf, time_out_buf diff --git a/legged_lab/runners/wmp_amp_runner.py b/legged_lab/runners/wmp_amp_runner.py index cdc0c1c..5763e47 100644 --- a/legged_lab/runners/wmp_amp_runner.py +++ b/legged_lab/runners/wmp_amp_runner.py @@ -124,6 +124,10 @@ def _build_amp(self): amp_cfg.get("replay_buffer_size", 100000), amp_cfg.get("grad_penalty_coef", 1.0), ) + min_std = self._make_min_action_std(amp_cfg.get("min_normalized_std")) + if min_std is not None and hasattr(self.alg, "set_min_action_std"): + self.alg.set_min_action_std(min_std) + print(f"[INFO] WMP-AMP min action std clamp: {min_std.detach().cpu().tolist()}") retarget_name = retarget_adapter.__class__.__name__ print( f"[INFO] WMP-AMP enabled: expert_dim={amp_data.observation_dim}, " @@ -143,6 +147,30 @@ def _make_amp_retarget_adapter(self, amp_cfg: dict): **retarget_kwargs, ) + def _make_min_action_std(self, min_normalized_std) -> torch.Tensor | None: + if min_normalized_std is None: + return None + min_normalized_std = torch.as_tensor(min_normalized_std, device=self.device, dtype=torch.float32) + if min_normalized_std.numel() != self.env.num_actions: + raise ValueError( + "amp.min_normalized_std must have one value per action: " + f"got {min_normalized_std.numel()}, expected {self.env.num_actions}." + ) + limits = getattr(self.env.robot.data, "soft_joint_pos_limits", None) + if limits is None: + limits = getattr(self.env.robot.data, "joint_pos_limits", None) + if limits is None: + raise RuntimeError("Cannot compute WMP min action std: robot joint position limits are unavailable.") + limits = limits[0].to(self.device) + action_range = limits[:, 1] - limits[:, 0] + if action_range.numel() != self.env.num_actions: + raise RuntimeError( + "Cannot compute WMP min action std: joint limit dim does not match action dim, " + f"limits={action_range.numel()}, actions={self.env.num_actions}." + ) + # 原版公式: min_std = min_normalized_std * (dof_upper_limit - dof_lower_limit)。 + return min_normalized_std * action_range + def _log_amp_joint_stats(self, amp_data: AMPLoader): if not hasattr(amp_data, "preloaded_s"): return From f75c90e935d6dd3390bf22be5aee1181ea6bf4cf Mon Sep 17 00:00:00 2001 From: m0_51228868 Date: Tue, 9 Jun 2026 09:40:29 +0800 Subject: [PATCH 3/6] fix WMP replay sampling semantics --- WMP_NETWORK_DESIGN_DIFF_REPORT.txt | 58 +++++++++---- legged_lab/world_models/wmp/models.py | 112 ++++++++++++++++++++++---- 2 files changed, 138 insertions(+), 32 deletions(-) diff --git a/WMP_NETWORK_DESIGN_DIFF_REPORT.txt b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt index fcdaefb..a4e0c0a 100644 --- a/WMP_NETWORK_DESIGN_DIFF_REPORT.txt +++ b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt @@ -808,8 +808,8 @@ critic 维度已经对齐，下一步不是改维度，而是验收字段语义: - 后续只需在真实 depth 分布异常时，再回到相机图像统计和可视化检查。 -12.5 World model replay 采样仍不是逐行等价 ------------------------------------------- +12.5 World model replay 采样复查与修正 +-------------------------------------- 原版 train_world_model(): - p = wm_dataset_size / sum(wm_dataset_size) 按 env episode 长度加权采样 env。 - batch_length = min(sampled_env_dataset_size.min(), config.batch_length)。 @@ -817,29 +817,56 @@ critic 维度已经对齐，下一步不是改维度，而是验收字段语义: - 非真实相机 env 的 image 在 sample 时即时由 DepthPredictor(forward_height_map, prop) 生成。 - is_first 每个 batch 第一帧置 1。 -当前 WMPFixedEpisodeReplayBuffer: - - 也按 dataset_size 加权采样 env，并把 batch 第一帧 is_first 置 1。 - - 但 replay 存储、episode finish、image 缺失标记、depth_index_slot 和 sample_real_depth 的实现不是原版逐行迁移。 - - 当前 train_start_steps 使用 replay.num_steps；原版使用 sum_wm_dataset_size。 +本轮复查发现的严重问题: + - WMPTrainingController 构造 WMPFixedEpisodeReplayBuffer 时传入 + depth_index_without_crawl_tilt、depth_index_inverse、depth_predictor_excluded_env_ids， + 但 WMPFixedEpisodeReplayBuffer.__init__ 之前没有接收这些参数。 + - 这不是采样分布差异，而是 WMP controller 初始化会直接 TypeError 的硬问题。 + +已修正: + - WMPFixedEpisodeReplayBuffer.__init__ 已接收并保存上述 depth index 参数。 + - camera_env_ids 不再排序去重，而是保留 env.depth_index 的传入顺序去重； + 这样真实 depth 的 camera slot 与 env id 映射不会被 replay buffer 打乱。 + - sample() 改回原版动态 batch_length 语义: + 先按 p_i = dataset_size_i / sum(dataset_size) 采样 env； + 再用 batch_length = min(sampled_dataset_size.min(), config.batch_length)； + 最后对每个 env 随机 end_idx 并切片 [end_idx-batch_length, end_idx)。 + - 因此短 episode / 失败 episode 不会因为长度小于 config.batch_length 被长期排除。 + - sample() 仍强制 out["is_first"][:, 0] = 1，并清零后续帧，和原版一致。 + - 非真实相机 env 的 image 仍先返回零占位，再由 _materialize_sampled_images() + 使用 DepthPredictor(forward_height_map, prop) 补齐，语义等价原版 sample 时即时生成。 + - 真实 depth 监督采样改为只使用 depth_index_without_crawl_tilt，并按 has_real_depth 过滤有效时间步。 影响: - 当前方向是语义接近，但仍应把 world model loss 曲线差异优先归因到 replay/tick/sample 细节，而不是 RSSM 网络结构。 + - 原先可能的严重问题是初始化直接失败，以及短 episode 被 world model 训练样本排除； + 这两项已经修正。 + - 现在 replay 采样不会跨 episode，dataset_size 之外的旧尾部数据不会被采到。 + - 仍需关注 world model loss 曲线差异，但优先级应从“采样硬 bug”降为“实现细节与时序差异”。 + +仍不是逐行等价: + - 当前使用 WMPFixedEpisodeReplayBuffer 封装 replay，内部结构不同于原版 wm_dataset/wm_buffer 两套 dict。 + - train_start_steps 当前使用 replay.num_steps；原版使用 sum_wm_dataset_size。两者在语义上接近， + 但如果后续引入容量裁剪或多 episode replay，可能需要重新对齐。 + - image 缺失标记使用 has_real_depth / camera slot 判断，不再依赖原版 np.where(env.depth_index == idx) 的逐行写法。 -12.6 DepthPredictor 监督采样仍有细节差异 ----------------------------------------- +12.6 DepthPredictor 监督采样已按原版主要语义收窄 +----------------------------------------------- 原版 train_depth_predictor(): - batch_idx 从 env.depth_index_without_crawl_tilt 采样。 - time_index 对每个 env 从 [0, wm_dataset_size[idx]] 采样。 - image 使用 env.depth_index_inverse[batch_idx] 映射到真实相机 image buffer。 当前: - - sample_real_depth() 从 camera_env_ids 中 dataset_size > 0 的 env 采样。 - - 是否严格排除 tilt/crawl 取决于 camera_env_ids、depth_index_without_crawl_tilt 和 has_real_depth 的联动。 + - sample_real_depth() 已从 depth_index_without_crawl_tilt 中采样 env。 + - depth_predictor_excluded_env_ids 会进一步排除 tilt/crawl env。 + - 采样时间步会按 has_real_depth 过滤，避免抽到没有真实 depth 写入的空 image。 + - image 通过 env_to_camera_slot / depth_index_inverse 语义映射到真实相机 image buffer。 风险: - 如果 tilt/crawl env 进入 DepthPredictor 监督采样，训练分布会比原版更难； - 如果 time_index 边界和原版不同，也会改变 early episode 的监督样本。 + - 原版 time_index 使用 np.random.randint(0, wm_dataset_size[idx] + 1)，当前使用有效真实 depth 时间步集合。 + 这样更能避免空 depth 样本，但不是逐行等价。 + - 如果真实 depth 更新频率或 has_real_depth 标记异常，DepthPredictor 监督样本数会下降。 12.7 Domain randomization 数值来源接近，但执行时机不同 @@ -896,9 +923,10 @@ critic 维度已经对齐，下一步不是改维度，而是验收字段语义: ------------------------ 如果下一步要继续缩小与原版训练效果的差距，建议优先级为: - 1. 对比一批 WMP replay sample: + 1. 抽样验收 WMP replay batch: prop/action/reward/is_first/image/forward_height_map 的 shape、时间索引、真实 depth 覆盖位置。 - 2. 对比 DepthPredictor 监督采样 env ids，确保默认只用 depth_index_without_crawl_tilt。 + 2. 记录 DepthPredictor 监督采样统计: + env ids 是否来自 depth_index_without_crawl_tilt、has_real_depth 命中率、loss 初始量级。 3. 做 PPO 接口 ablation: obs normalizer、rsl_rl v5 storage/update 细节。 4. 记录 reset cause 分布、episode length 分布、AMP policy/expert score 分布。 diff --git a/legged_lab/world_models/wmp/models.py b/legged_lab/world_models/wmp/models.py index 7167d2b..defd12e 100644 --- a/legged_lab/world_models/wmp/models.py +++ b/legged_lab/world_models/wmp/models.py @@ -235,17 +235,31 @@ def __init__( max_episode_steps: int, camera_env_ids: torch.Tensor, device: str, + depth_index_without_crawl_tilt: torch.Tensor | None = None, + depth_index_inverse: torch.Tensor | None = None, + depth_predictor_excluded_env_ids: torch.Tensor | None = None, ): self.num_envs = int(num_envs) self.max_episode_steps = int(max_episode_steps) self.device = torch.device(device) - self.camera_env_ids = camera_env_ids.detach().cpu().long().unique(sorted=True) - self.camera_env_ids = self.camera_env_ids[ - (self.camera_env_ids >= 0) & (self.camera_env_ids < self.num_envs) - ] + # 保留 env.depth_index 的顺序，避免真实 depth 的 camera slot 与 env id 映射错位。 + self.camera_env_ids = self._ordered_unique_valid(camera_env_ids) self.env_to_camera_slot = torch.full((self.num_envs,), -1, dtype=torch.long) if self.camera_env_ids.numel() > 0: self.env_to_camera_slot[self.camera_env_ids] = torch.arange(self.camera_env_ids.numel(), dtype=torch.long) + self.depth_index_inverse = self._make_depth_index_inverse(depth_index_inverse) + if depth_index_without_crawl_tilt is None: + depth_index_without_crawl_tilt = self.camera_env_ids + self.depth_index_without_crawl_tilt = self._ordered_unique_valid(depth_index_without_crawl_tilt) + if depth_predictor_excluded_env_ids is not None: + excluded = torch.zeros(self.num_envs, dtype=torch.bool) + excluded[self._ordered_unique_valid(depth_predictor_excluded_env_ids)] = True + self.depth_index_without_crawl_tilt = self.depth_index_without_crawl_tilt[ + ~excluded[self.depth_index_without_crawl_tilt] + ] + if self.depth_index_without_crawl_tilt.numel() > 0: + has_camera = self.env_to_camera_slot[self.depth_index_without_crawl_tilt] >= 0 + self.depth_index_without_crawl_tilt = self.depth_index_without_crawl_tilt[has_camera] self.current_index = torch.zeros(self.num_envs, dtype=torch.long) self.dataset_size = torch.zeros(self.num_envs, dtype=torch.long) @@ -254,6 +268,7 @@ def __init__( self._dataset: dict[str, torch.Tensor] | None = None self._current_image: torch.Tensor | None = None self._dataset_image: torch.Tensor | None = None + self._image_shape: tuple[int, ...] | None = None def __len__(self): return int(self.has_episode.sum().item()) @@ -305,36 +320,40 @@ def finish_episode(self, env_id: int): self.has_episode[env_id] = True def can_sample(self, batch_size: int, batch_length: int) -> bool: + del batch_length if batch_size <= 0 or self._dataset is None: return False - return bool(torch.any(self.dataset_size >= int(batch_length)).item()) + return bool(torch.any(self.dataset_size > 1).item()) def can_sample_real_depth(self, batch_size: int) -> bool: if batch_size <= 0 or self._dataset is None or self._dataset_image is None: return False - if self.camera_env_ids.numel() == 0: - return False - valid_camera_envs = self.camera_env_ids[self.dataset_size[self.camera_env_ids] > 0] - return bool(valid_camera_envs.numel() > 0) + valid_envs, _ = self._valid_real_depth_envs_and_counts() + return bool(valid_envs.numel() > 0) def sample(self, batch_size: int, batch_length: int) -> dict[str, torch.Tensor]: if not self.can_sample(batch_size, batch_length): raise RuntimeError("WMP fixed replay buffer does not have enough sequence data.") assert self._dataset is not None - batch_length = int(batch_length) - valid_envs = (self.dataset_size >= batch_length).nonzero(as_tuple=False).flatten() + requested_batch_length = int(batch_length) + valid_envs = (self.dataset_size > 1).nonzero(as_tuple=False).flatten() lengths = self.dataset_size[valid_envs].float() + # 原版 WMP 采样概率: p_i = dataset_size_i / sum_j(dataset_size_j)。 probs = lengths / lengths.sum().clamp_min(1.0) sampled_slots = torch.multinomial(probs, int(batch_size), replacement=True) env_ids = valid_envs[sampled_slots] + batch_length = min(int(self.dataset_size[env_ids].min().item()), requested_batch_length) + if batch_length <= 1: + raise RuntimeError("WMP fixed replay sampled sequence length must be greater than 1.") batch = {key: [] for key in self._dataset.keys()} batch["image"] = [] for env_id_tensor in env_ids: env_id = int(env_id_tensor.item()) ep_len = int(self.dataset_size[env_id].item()) - start = int(torch.randint(0, ep_len - batch_length + 1, (1,)).item()) + end = int(torch.randint(batch_length, ep_len + 1, (1,)).item()) + start = end - batch_length for key, value in self._dataset.items(): batch[key].append(value[env_id, start : start + batch_length]) batch["image"].append(self._sample_image(env_id, start, batch_length)) @@ -351,9 +370,8 @@ def sample_real_depth(self, batch_size: int) -> dict[str, torch.Tensor]: assert self._dataset is not None assert self._dataset_image is not None - valid_envs = self.camera_env_ids[self.dataset_size[self.camera_env_ids] > 0] - lengths = self.dataset_size[valid_envs].float() - probs = lengths / lengths.sum().clamp_min(1.0) + valid_envs, valid_counts = self._valid_real_depth_envs_and_counts() + probs = valid_counts.float() / valid_counts.float().sum().clamp_min(1.0) sampled_slots = torch.multinomial(probs, int(batch_size), replacement=True) env_ids = valid_envs[sampled_slots] @@ -361,7 +379,7 @@ def sample_real_depth(self, batch_size: int) -> dict[str, torch.Tensor]: for env_id_tensor in env_ids: env_id = int(env_id_tensor.item()) ep_len = int(self.dataset_size[env_id].item()) - step_id = int(torch.randint(0, ep_len, (1,)).item()) + step_id = self._sample_real_depth_step(env_id, ep_len) camera_slot = int(self.env_to_camera_slot[env_id].item()) batch["forward_height_map"].append(self._dataset["forward_height_map"][env_id, step_id]) batch["prop"].append(self._dataset["prop"][env_id, step_id]) @@ -379,6 +397,7 @@ def _init_storage(self, transition: dict[str, torch.Tensor]): self._dataset[key] = torch.zeros(shape, device=self.device, dtype=torch.float32) image_shape = tuple(transition["image"].shape) + self._image_shape = image_shape camera_count = int(self.camera_env_ids.numel()) if camera_count > 0: shape = (camera_count, self.max_episode_steps) + image_shape @@ -387,8 +406,67 @@ def _init_storage(self, transition: dict[str, torch.Tensor]): def _sample_image(self, env_id: int, start: int, batch_length: int) -> torch.Tensor: if self._dataset_image is None: - raise RuntimeError("WMP fixed replay image storage has not been initialized.") + if self._image_shape is None: + raise RuntimeError("WMP fixed replay image storage has not been initialized.") + return torch.zeros((batch_length,) + self._image_shape, device=self.device, dtype=torch.float32) camera_slot = int(self.env_to_camera_slot[env_id].item()) if camera_slot >= 0: return self._dataset_image[camera_slot, start : start + batch_length] return torch.zeros_like(self._dataset_image[0, start : start + batch_length]) + + def _ordered_unique_valid(self, ids: torch.Tensor | None) -> torch.Tensor: + if ids is None: + return torch.empty(0, dtype=torch.long) + ids = torch.as_tensor(ids, device="cpu", dtype=torch.long).flatten() + keep = torch.zeros(self.num_envs, dtype=torch.bool) + out = [] + for env_id in ids.tolist(): + env_id = int(env_id) + if 0 <= env_id < self.num_envs and not bool(keep[env_id].item()): + keep[env_id] = True + out.append(env_id) + return torch.tensor(out, dtype=torch.long) + + def _make_depth_index_inverse(self, depth_index_inverse: torch.Tensor | None) -> torch.Tensor: + if depth_index_inverse is not None: + inverse = torch.as_tensor(depth_index_inverse, device="cpu", dtype=torch.long).flatten() + if inverse.numel() == self.num_envs: + return inverse.clone() + inverse = torch.full((self.num_envs,), -1, dtype=torch.long) + if self.camera_env_ids.numel() > 0: + inverse[self.camera_env_ids] = torch.arange(self.camera_env_ids.numel(), dtype=torch.long) + return inverse + + def _valid_real_depth_envs_and_counts(self) -> tuple[torch.Tensor, torch.Tensor]: + if self._dataset is None or self.depth_index_without_crawl_tilt.numel() == 0: + empty = torch.empty(0, dtype=torch.long) + return empty, empty + has_real_depth = self._dataset.get("has_real_depth") + valid_envs = [] + valid_counts = [] + for env_id in self.depth_index_without_crawl_tilt.tolist(): + env_id = int(env_id) + ep_len = int(self.dataset_size[env_id].item()) + if ep_len <= 0: + continue + if has_real_depth is None: + valid_count = ep_len + else: + valid_mask = has_real_depth[env_id, :ep_len].reshape(ep_len, -1)[:, 0] > 0.5 + valid_count = int(valid_mask.sum().item()) + if valid_count > 0: + valid_envs.append(env_id) + valid_counts.append(valid_count) + return torch.tensor(valid_envs, dtype=torch.long), torch.tensor(valid_counts, dtype=torch.long) + + def _sample_real_depth_step(self, env_id: int, ep_len: int) -> int: + has_real_depth = self._dataset.get("has_real_depth") if self._dataset is not None else None + if has_real_depth is None: + return int(torch.randint(0, ep_len, (1,)).item()) + valid_steps = (has_real_depth[env_id, :ep_len].reshape(ep_len, -1)[:, 0] > 0.5).nonzero( + as_tuple=False + ).flatten() + if valid_steps.numel() == 0: + return int(torch.randint(0, ep_len, (1,)).item()) + slot = int(torch.randint(0, valid_steps.numel(), (1,)).item()) + return int(valid_steps[slot].item()) From 62098e87a2613a4b8c1a5936f9d7e55fa87142dc Mon Sep 17 00:00:00 2001 From: m0_51228868 Date: Tue, 9 Jun 2026 10:02:58 +0800 Subject: [PATCH 4/6] align WMP domain randomization semantics --- WMP_NETWORK_DESIGN_DIFF_REPORT.txt | 40 +++- legged_lab/envs/a1/a1_config.py | 2 +- legged_lab/mdp/__init__.py | 1 + legged_lab/mdp/events.py | 355 +++++++++++++++++++++++++++++ 4 files changed, 385 insertions(+), 13 deletions(-) create mode 100644 legged_lab/mdp/events.py diff --git a/WMP_NETWORK_DESIGN_DIFF_REPORT.txt b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt index a4e0c0a..a7f04d5 100644 --- a/WMP_NETWORK_DESIGN_DIFF_REPORT.txt +++ b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt @@ -869,20 +869,36 @@ critic 维度已经对齐，下一步不是改维度，而是验收字段语义: - 如果真实 depth 更新频率或 has_real_depth 标记异常，DepthPredictor 监督样本数会下降。 -12.7 Domain randomization 数值来源接近，但执行时机不同 ------------------------------------------------------- +12.7 Domain randomization 内部自洽已补齐，但最终物理效果不能保证完全一致 +----------------------------------------------------------------------- 原版: - - friction/restitution/base mass/link mass/COM/gains 在 IsaacGym 创建 env 或 reset callback 中随机。 - - motor_strength 在每个 sim substep 对 torque 乘随机系数。 - - action_latency 每次 step 随机整数延迟。 + - friction/restitution 在 env 创建时为每个 env 采样标量，并写入该 env 的 rigid shapes。 + - base mass、link mass、COM 在 IsaacGym 创建 env 时随机。 + - gains 在每次 reset 中重新采样 randomized_p_gains/randomized_d_gains。 + - motor_strength 在每个 sim substep 对最终 torque 乘随机系数。 + - action_latency 每次 step 随机整数延迟；延迟时使用 last_actions 重新计算 torque。 + +当前已修订: + - 新增 legged_lab.mdp.wmp_recording_* EventTerm，friction/restitution/base mass/COM/gains 的采样值会同时写入物理仿真和 WMP critic privileged buffer。 + - A1 WMP 的数值范围已与原版主要配置对齐: + friction [0.5, 2.0]，restitution [0.0, 0.0]， + base added mass [0.0, 3.0]，link mass scale [0.8, 1.2]， + COM xyz [-0.05, 0.05]，PD gains scale [0.8, 1.2]， + motor_strength [0.8, 1.2]，latency 0 到约 0.005s，push interval 15s/max xy vel 1.0。 + - gains EventTerm 已从 startup 改为 reset，和原版每个 episode reset 重采 gains 的语义一致。 + - recording material EventTerm 按 env 采样 friction/restitution 标量并写入所有相关 shapes，critic 读取同一采样值，不再依赖独立 fallback sample。 + - recording COM EventTerm 记录量为 new_com - default_com；critic 中再乘 obs_scales.com_pos，语义对应原版 randomized_com_pos * obs_scales.com_pos。 + - recording gains EventTerm 记录量为 randomized_gain / default_gain - 1；critic 中再乘 obs_scales.pd_gains，语义对应原版。 + +仍不等价: + - motor_strength 仍在 reset 时采样，并作用于 action target/position target；原版是在每个 sim substep 对已经算出的 torque 重新采样乘法系数。 + - action_delay 当前 sim_step 模式会在 substep 级选择 last_processed_actions 或当前 processed_actions；原版是在 torque 计算入口选择 last_actions 或 actions。 + - IsaacLab/PhysX 与原版 IsaacGym 的材质、质量、COM、显式 actuator 计算路径不同；即使同一范围和同一记录量，也不能保证完全相同的接触/力矩轨迹。 -当前: - - friction/restitution/base mass/COM/gains 使用 IsaacLab EventTerm 或记录型 EventTerm。 - - motor_strength 在 reset 时采样，并作用于 action target。 - - action_delay 使用 IsaacLab DelayBuffer / sim_step 语义。 - -影响: - critic privileged buffer 已尽量记录对应随机量，但物理执行时机、作用位置和随机重采样频率仍不等价。 +结论: + 当前可以保证 domain randomization 的 critic 输入与本项目实际 EventTerm 随机量内部自洽； + 不能保证最终训练效果与原版 WMP 完全一致。若要进一步逼近，优先项是实现 torque-level、 + per-substep 的 motor_strength hook，并对 action latency 在 torque 入口处做等价延迟。 12.8 Reward / command 基本对齐，但仍需跑统计 diff --git a/legged_lab/envs/a1/a1_config.py b/legged_lab/envs/a1/a1_config.py index fcdb528..c6b7814 100644 --- a/legged_lab/envs/a1/a1_config.py +++ b/legged_lab/envs/a1/a1_config.py @@ -297,7 +297,7 @@ def __post_init__(self): ) self.domain_rand.events.randomize_actuator_gains = EventTerm( func=mdp.wmp_recording_randomize_actuator_gains, - mode="startup", + mode="reset", params={ "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), "stiffness_distribution_params": (0.8, 1.2), diff --git a/legged_lab/mdp/__init__.py b/legged_lab/mdp/__init__.py index 584d731..cb57058 100644 --- a/legged_lab/mdp/__init__.py +++ b/legged_lab/mdp/__init__.py @@ -13,4 +13,5 @@ from isaaclab.envs.mdp import * +from .events import * from .rewards import * diff --git a/legged_lab/mdp/events.py b/legged_lab/mdp/events.py new file mode 100644 index 0000000..c0cf594 --- /dev/null +++ b/legged_lab/mdp/events.py @@ -0,0 +1,355 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# Original code is licensed under BSD-3-Clause. +# +# Copyright (c) 2025-2026, The Legged Lab Project Developers. +# All rights reserved. +# Modifications are licensed under BSD-3-Clause. +# +# This file contains code derived from Isaac Lab Project (BSD-3-Clause license) +# with modifications by Legged Lab Project (BSD-3-Clause license). + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING, Literal + +import torch + +import isaaclab.utils.math as math_utils +from isaaclab.actuators import ImplicitActuator +from isaaclab.assets import Articulation, RigidObject +from isaaclab.envs.mdp import events as isaac_events +from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + +def _env_ids_cpu(env: ManagerBasedEnv, env_ids: Sequence[int] | torch.Tensor | slice | None) -> torch.Tensor: + if env_ids is None or isinstance(env_ids, slice): + return torch.arange(env.scene.num_envs, device="cpu") + if isinstance(env_ids, torch.Tensor): + return env_ids.detach().cpu() + return torch.tensor(env_ids, dtype=torch.long, device="cpu") + + +def _env_ids_device(env: ManagerBasedEnv, env_ids_cpu: torch.Tensor) -> torch.Tensor: + return env_ids_cpu.to(device=env.device, dtype=torch.long) + + +def _sample_distribution( + distribution_parameters: tuple[float | torch.Tensor, float | torch.Tensor], + shape: tuple[int, ...], + device: str | torch.device, + distribution: Literal["uniform", "log_uniform", "gaussian"], +) -> torch.Tensor: + if distribution == "uniform": + return math_utils.sample_uniform(*distribution_parameters, shape, device=device) + if distribution == "log_uniform": + return math_utils.sample_log_uniform(*distribution_parameters, shape, device=device) + if distribution == "gaussian": + return math_utils.sample_gaussian(*distribution_parameters, shape, device=device) + raise NotImplementedError(f"Unknown randomization distribution: '{distribution}'.") + + +def _resolve_body_ids(asset: RigidObject | Articulation, asset_cfg: SceneEntityCfg) -> torch.Tensor: + if asset_cfg.body_ids == slice(None): + return torch.arange(asset.num_bodies, dtype=torch.int, device="cpu") + return torch.tensor(asset_cfg.body_ids, dtype=torch.int, device="cpu") + + +class wmp_recording_randomize_rigid_body_material(isaac_events.randomize_rigid_body_material): + """Randomize material properties and record WMP critic privileged values. + + 原版 WMP 在 env 创建时为每个 env 采样一个 friction/restitution 标量，并把同一个值 + 写到该 env 的 rigid shapes；critic 也读取同一个标量。这里保持同源采样： + `wmp_priv_friction = sampled_static_friction`, + `wmp_priv_restitution = sampled_restitution`。 + """ + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + static_friction_range: tuple[float, float], + dynamic_friction_range: tuple[float, float], + restitution_range: tuple[float, float], + num_buckets: int, + asset_cfg: SceneEntityCfg, + make_consistent: bool = False, + ): + env_ids_cpu = _env_ids_cpu(env, env_ids) + env_ids_device = _env_ids_device(env, env_ids_cpu) + total_num_shapes = self.asset.root_physx_view.max_shapes + + static_friction = math_utils.sample_uniform( + static_friction_range[0], static_friction_range[1], (len(env_ids_cpu), 1, 1), device="cpu" + ) + dynamic_friction = math_utils.sample_uniform( + dynamic_friction_range[0], dynamic_friction_range[1], (len(env_ids_cpu), 1, 1), device="cpu" + ) + if make_consistent: + dynamic_friction = torch.minimum(static_friction, dynamic_friction) + restitution = math_utils.sample_uniform( + restitution_range[0], restitution_range[1], (len(env_ids_cpu), 1, 1), device="cpu" + ) + material_samples = torch.cat([static_friction, dynamic_friction, restitution], dim=-1).repeat( + 1, total_num_shapes, 1 + ) + + materials = self.asset.root_physx_view.get_material_properties() + if self.num_shapes_per_body is not None: + for body_id in self.asset_cfg.body_ids: + start_idx = sum(self.num_shapes_per_body[:body_id]) + end_idx = start_idx + self.num_shapes_per_body[body_id] + materials[env_ids_cpu, start_idx:end_idx] = material_samples[:, start_idx:end_idx] + else: + materials[env_ids_cpu] = material_samples + self.asset.root_physx_view.set_material_properties(materials, env_ids_cpu) + + if hasattr(env, "wmp_priv_friction"): + env.wmp_priv_friction[env_ids_device] = static_friction[:, 0].to(env.device) + if hasattr(env, "wmp_priv_restitution"): + env.wmp_priv_restitution[env_ids_device] = restitution[:, 0].to(env.device) + + +class wmp_recording_randomize_rigid_body_mass(ManagerTermBase): + """Randomize body mass and record the additive base-mass term used by WMP critic.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.asset: RigidObject | Articulation = env.scene[self.asset_cfg.name] + + operation = cfg.params["operation"] + if operation == "scale": + if "mass_distribution_params" in cfg.params: + isaac_events._validate_scale_range( + cfg.params["mass_distribution_params"], "mass_distribution_params", allow_zero=False + ) + elif operation not in ("abs", "add"): + raise ValueError(f"WMP mass randomization does not support operation: '{operation}'.") + if cfg.params.get("min_mass") is not None and cfg.params.get("min_mass") < 1.0e-6: + raise ValueError("WMP mass randomization requires 'min_mass' >= 1e-6.") + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + asset_cfg: SceneEntityCfg, + mass_distribution_params: tuple[float, float], + operation: Literal["add", "scale", "abs"], + distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", + recompute_inertia: bool = True, + min_mass: float = 1.0e-6, + ): + env_ids_cpu = _env_ids_cpu(env, env_ids) + env_ids_device = _env_ids_device(env, env_ids_cpu) + body_ids = _resolve_body_ids(self.asset, self.asset_cfg) + + masses = self.asset.root_physx_view.get_masses() + default_mass = self.asset.data.default_mass[env_ids_cpu[:, None], body_ids].clone() + masses[env_ids_cpu[:, None], body_ids] = default_mass + + samples = _sample_distribution(mass_distribution_params, (len(env_ids_cpu), 1), masses.device, distribution) + if operation == "add": + randomized_mass = default_mass + samples + recorded_added_mass = samples + elif operation == "scale": + randomized_mass = default_mass * samples + recorded_added_mass = randomized_mass[:, :1] - default_mass[:, :1] + elif operation == "abs": + randomized_mass = samples.expand_as(default_mass) + recorded_added_mass = randomized_mass[:, :1] - default_mass[:, :1] + else: + raise NotImplementedError(f"Unknown mass randomization operation: '{operation}'.") + + masses[env_ids_cpu[:, None], body_ids] = torch.clamp(randomized_mass, min=min_mass) + self.asset.root_physx_view.set_masses(masses, env_ids_cpu) + + if recompute_inertia: + ratios = masses[env_ids_cpu[:, None], body_ids] / self.asset.data.default_mass[env_ids_cpu[:, None], body_ids] + inertias = self.asset.root_physx_view.get_inertias() + if isinstance(self.asset, Articulation): + inertias[env_ids_cpu[:, None], body_ids] = ( + self.asset.data.default_inertia[env_ids_cpu[:, None], body_ids] * ratios[..., None] + ) + else: + inertias[env_ids_cpu] = self.asset.data.default_inertia[env_ids_cpu] * ratios + self.asset.root_physx_view.set_inertias(inertias, env_ids_cpu) + + if hasattr(env, "wmp_priv_added_mass"): + env.wmp_priv_added_mass[env_ids_device] = recorded_added_mass.to(env.device) + + +def wmp_recording_randomize_rigid_body_com( + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + com_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg, +): + """Randomize CoM offset and record the sampled offset used by WMP critic. + + 记录量为 `new_com - default_com`，与原版 critic 中 `randomized_com_pos * obs_scales.com_pos` + 的语义一致。 + """ + + asset: Articulation = env.scene[asset_cfg.name] + env_ids_cpu = _env_ids_cpu(env, env_ids) + env_ids_device = _env_ids_device(env, env_ids_cpu) + body_ids = _resolve_body_ids(asset, asset_cfg) + + ranges = torch.tensor([com_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]], device="cpu") + offsets = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids_cpu), 3), device="cpu") + + coms = asset.root_physx_view.get_coms().clone() + if hasattr(env, "wmp_priv_default_coms"): + default_coms = env.wmp_priv_default_coms.detach().cpu() + coms[env_ids_cpu[:, None], body_ids, :3] = default_coms[env_ids_cpu[:, None], body_ids, :3] + coms[env_ids_cpu[:, None], body_ids, :3] += offsets.unsqueeze(1) + asset.root_physx_view.set_coms(coms, env_ids_cpu) + + if hasattr(env, "wmp_priv_com_pos"): + env.wmp_priv_com_pos[env_ids_device] = offsets.to(env.device) + + +class wmp_recording_randomize_actuator_gains(ManagerTermBase): + """Randomize actuator gains and record `(randomized_gain / default_gain - 1)`. + + 原版 WMP critic 使用: + p_gain_scale = randomized_p_gains / p_gains - 1 + d_gain_scale = randomized_d_gains / d_gains - 1 + BaseEnv 后续会再乘 `obs_scales.pd_gains`。 + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.asset: RigidObject | Articulation = env.scene[self.asset_cfg.name] + + operation = cfg.params["operation"] + if operation == "scale": + if "stiffness_distribution_params" in cfg.params: + isaac_events._validate_scale_range( + cfg.params["stiffness_distribution_params"], "stiffness_distribution_params", allow_zero=False + ) + if "damping_distribution_params" in cfg.params: + isaac_events._validate_scale_range( + cfg.params["damping_distribution_params"], "damping_distribution_params" + ) + elif operation not in ("abs", "add"): + raise ValueError(f"WMP actuator gain randomization does not support operation: '{operation}'.") + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + asset_cfg: SceneEntityCfg, + stiffness_distribution_params: tuple[float, float] | None = None, + damping_distribution_params: tuple[float, float] | None = None, + operation: Literal["add", "scale", "abs"] = "abs", + distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", + ): + if env_ids is None or isinstance(env_ids, slice): + env_ids_asset = torch.arange(env.scene.num_envs, device=self.asset.device) + else: + env_ids_asset = env_ids.to(self.asset.device) + env_ids_device = env_ids_asset.to(env.device) + + for actuator in self.asset.actuators.values(): + if isinstance(self.asset_cfg.joint_ids, slice): + actuator_indices = slice(None) + if isinstance(actuator.joint_indices, slice): + global_indices = slice(None) + elif isinstance(actuator.joint_indices, torch.Tensor): + global_indices = actuator.joint_indices.to(self.asset.device) + else: + raise TypeError("Actuator joint indices must be a slice or a torch.Tensor.") + elif isinstance(actuator.joint_indices, slice): + global_indices = actuator_indices = torch.tensor(self.asset_cfg.joint_ids, device=self.asset.device) + else: + actuator_joint_indices = actuator.joint_indices + asset_joint_ids = torch.tensor(self.asset_cfg.joint_ids, device=self.asset.device) + actuator_indices = torch.nonzero(torch.isin(actuator_joint_indices, asset_joint_ids)).view(-1) + if len(actuator_indices) == 0: + continue + global_indices = actuator_joint_indices[actuator_indices] + + if stiffness_distribution_params is not None: + stiffness = actuator.stiffness[env_ids_asset].clone() + default_stiffness = self.asset.data.default_joint_stiffness[env_ids_asset][:, global_indices].clone() + stiffness[:, actuator_indices] = self._randomize_gain( + default_stiffness, stiffness_distribution_params, operation, distribution + ) + actuator.stiffness[env_ids_asset] = stiffness + if hasattr(env, "wmp_priv_p_gain_scale"): + self._record_gain_scale( + env.wmp_priv_p_gain_scale, + env_ids_device, + global_indices, + default_stiffness, + stiffness[:, actuator_indices], + ) + if isinstance(actuator, ImplicitActuator): + self.asset.write_joint_stiffness_to_sim( + stiffness, joint_ids=actuator.joint_indices, env_ids=env_ids_asset + ) + + if damping_distribution_params is not None: + damping = actuator.damping[env_ids_asset].clone() + default_damping = self.asset.data.default_joint_damping[env_ids_asset][:, global_indices].clone() + damping[:, actuator_indices] = self._randomize_gain( + default_damping, damping_distribution_params, operation, distribution + ) + actuator.damping[env_ids_asset] = damping + if hasattr(env, "wmp_priv_d_gain_scale"): + self._record_gain_scale( + env.wmp_priv_d_gain_scale, + env_ids_device, + global_indices, + default_damping, + damping[:, actuator_indices], + ) + if isinstance(actuator, ImplicitActuator): + self.asset.write_joint_damping_to_sim( + damping, joint_ids=actuator.joint_indices, env_ids=env_ids_asset + ) + + def _randomize_gain( + self, + default_gain: torch.Tensor, + distribution_parameters: tuple[float, float], + operation: Literal["add", "scale", "abs"], + distribution: Literal["uniform", "log_uniform", "gaussian"], + ) -> torch.Tensor: + samples = _sample_distribution(distribution_parameters, default_gain.shape, default_gain.device, distribution) + if operation == "add": + return default_gain + samples + if operation == "scale": + return default_gain * samples + if operation == "abs": + return samples + raise NotImplementedError(f"Unknown actuator gain randomization operation: '{operation}'.") + + def _record_gain_scale( + self, + target: torch.Tensor, + env_ids: torch.Tensor, + global_indices: torch.Tensor | slice, + default_gain: torch.Tensor, + randomized_gain: torch.Tensor, + ): + gain_scale = randomized_gain / torch.clamp(default_gain, min=1.0e-8) - 1.0 + if isinstance(global_indices, slice): + target[env_ids] = gain_scale.to(target.device) + else: + target[env_ids[:, None], global_indices.to(target.device)] = gain_scale.to(target.device) + + +__all__ = [ + "wmp_recording_randomize_rigid_body_material", + "wmp_recording_randomize_rigid_body_mass", + "wmp_recording_randomize_rigid_body_com", + "wmp_recording_randomize_actuator_gains", +] From 0797083666c8a525391b15256106871e56b05a02 Mon Sep 17 00:00:00 2001 From: m0_51228868 Date: Tue, 9 Jun 2026 10:15:37 +0800 Subject: [PATCH 5/6] move WMP motor strength to torque path --- WMP_NETWORK_DESIGN_DIFF_REPORT.txt | 17 +++++++++------ legged_lab/envs/base/base_env.py | 33 +++++++++++++++++------------- 2 files changed, 30 insertions(+), 20 deletions(-) diff --git a/WMP_NETWORK_DESIGN_DIFF_REPORT.txt b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt index a7f04d5..1a46c9e 100644 --- a/WMP_NETWORK_DESIGN_DIFF_REPORT.txt +++ b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt @@ -869,8 +869,8 @@ critic 维度已经对齐，下一步不是改维度，而是验收字段语义: - 如果真实 depth 更新频率或 has_real_depth 标记异常，DepthPredictor 监督样本数会下降。 -12.7 Domain randomization 内部自洽已补齐，但最终物理效果不能保证完全一致 ------------------------------------------------------------------------ +12.7 Domain randomization 内部自洽已补齐，motor_strength 已改为 torque-level +------------------------------------------------------------------------ 原版: - friction/restitution 在 env 创建时为每个 env 采样标量，并写入该 env 的 rigid shapes。 - base mass、link mass、COM 在 IsaacGym 创建 env 时随机。 @@ -889,16 +889,21 @@ critic 维度已经对齐，下一步不是改维度，而是验收字段语义: - recording material EventTerm 按 env 采样 friction/restitution 标量并写入所有相关 shapes，critic 读取同一采样值，不再依赖独立 fallback sample。 - recording COM EventTerm 记录量为 new_com - default_com；critic 中再乘 obs_scales.com_pos，语义对应原版 randomized_com_pos * obs_scales.com_pos。 - recording gains EventTerm 记录量为 randomized_gain / default_gain - 1；critic 中再乘 obs_scales.pd_gains，语义对应原版。 + - motor_strength 不再乘 action/position target；现在每个 physics substep 在 IsaacLab actuator 生成 torque 后采样 s， + 执行 tau_final = s * tau，并在 sim.step 前覆盖 PhysX dof actuation force。这个插入位置与原版 + _compute_torques() 后、set_dof_actuation_force_tensor() 前的语义一致。 仍不等价: - - motor_strength 仍在 reset 时采样，并作用于 action target/position target；原版是在每个 sim substep 对已经算出的 torque 重新采样乘法系数。 - action_delay 当前 sim_step 模式会在 substep 级选择 last_processed_actions 或当前 processed_actions；原版是在 torque 计算入口选择 last_actions 或 actions。 - - IsaacLab/PhysX 与原版 IsaacGym 的材质、质量、COM、显式 actuator 计算路径不同；即使同一范围和同一记录量，也不能保证完全相同的接触/力矩轨迹。 + - motor_strength hook 依赖 IsaacLab Articulation 的 torque buffer；A1 WMP 使用显式 DCMotor 路径时语义接近原版， + 但这不是 IsaacLab 公开跨 actuator 的通用接口。 + - IsaacLab/PhysX 与原版 IsaacGym 的材质、质量、COM、显式 actuator 计算路径不同；即使同一范围、同一记录量和同一 torque-level s， + 也不能保证完全相同的接触/力矩轨迹。 结论: 当前可以保证 domain randomization 的 critic 输入与本项目实际 EventTerm 随机量内部自洽； - 不能保证最终训练效果与原版 WMP 完全一致。若要进一步逼近，优先项是实现 torque-level、 - per-substep 的 motor_strength hook，并对 action latency 在 torque 入口处做等价延迟。 + motor_strength 的插入位置与每 substep 重采样频率已按原版主要语义对齐。整体仍不能保证最终训练效果 + 与原版 WMP 完全一致，剩余优先项是 action latency 的 torque 入口等价延迟和 IsaacGym/IsaacLab backend 差异验收。 12.8 Reward / command 基本对齐，但仍需跑统计 diff --git a/legged_lab/envs/base/base_env.py b/legged_lab/envs/base/base_env.py index 7adad11..e6bb072 100644 --- a/legged_lab/envs/base/base_env.py +++ b/legged_lab/envs/base/base_env.py @@ -332,8 +332,6 @@ def init_buffers(self): self.action_buffer.set_time_lag(time_lags, torch.arange(self.num_envs, device=self.device)) self.last_processed_actions = self.robot.data.default_joint_pos.clone() self.motor_strength = torch.ones(self.num_envs, self.num_actions, device=self.device) - if self.cfg.domain_rand.motor_strength.enable: - self._randomize_motor_strength(torch.arange(self.num_envs, device=self.device)) self.robot_cfg = SceneEntityCfg(name="robot") self.robot_cfg.resolve(self.scene) @@ -452,8 +450,6 @@ def reset(self, env_ids): self.actor_obs_buffer.reset(env_ids) self.critic_obs_buffer.reset(env_ids) self.action_buffer.reset(env_ids) - if self.cfg.domain_rand.motor_strength.enable: - self._randomize_motor_strength(env_ids) if hasattr(self, "last_push_step_buf"): self.last_push_step_buf[env_ids] = -1 if hasattr(self, "push_recovered_buf"): @@ -468,8 +464,6 @@ def step(self, actions: torch.Tensor): delayed_actions = self.action_buffer.compute(actions) cliped_actions = torch.clip(delayed_actions, -self.clip_actions, self.clip_actions).to(self.device) - if self.cfg.domain_rand.motor_strength.enable: - cliped_actions = cliped_actions * self.motor_strength processed_actions = cliped_actions * self.action_scale + self.robot.data.default_joint_pos latency_steps = 0 if self.cfg.domain_rand.action_delay.enable and self.action_delay_mode == "sim_step": @@ -488,6 +482,8 @@ def step(self, actions: torch.Tensor): action_target = self.last_processed_actions if substep < latency_steps else processed_actions self.robot.set_joint_position_target(action_target) self.scene.write_data_to_sim() + if self.cfg.domain_rand.motor_strength.enable: + self._apply_motor_strength_to_torque() self.sim.step(render=False) if self.sim_step_counter % self.render_interval == 0 and is_rendering: self.sim.render() @@ -1017,16 +1013,25 @@ def _post_process_rewards(self, reward_buf: torch.Tensor) -> torch.Tensor: reward_buf = torch.clip(reward_buf, min=0.0) return reward_buf - def _randomize_motor_strength(self, env_ids: torch.Tensor): - if len(env_ids) == 0: - return + def _apply_motor_strength_to_torque(self): low, high = self.cfg.domain_rand.motor_strength.range + effort_target = getattr(self.robot, "_joint_effort_target_sim", None) + all_indices = getattr(self.robot, "_ALL_INDICES", None) + if effort_target is None or all_indices is None: + raise RuntimeError("WMP motor_strength requires IsaacLab Articulation torque buffers.") + if self.motor_strength.shape != effort_target.shape: + self.motor_strength = torch.ones_like(effort_target) + if abs(float(high) - float(low)) < 1.0e-8: - self.motor_strength[env_ids] = float(low) - return - self.motor_strength[env_ids] = torch.empty( - len(env_ids), self.num_actions, device=self.device - ).uniform_(float(low), float(high)) + self.motor_strength.fill_(float(low)) + else: + self.motor_strength.uniform_(float(low), float(high)) + + # 原版 WMP 在 _compute_torques() 完成 torque clip 后执行 tau_final = s * tau。 + # 这里等待 IsaacLab actuator 生成 _joint_effort_target_sim 后再乘 s，并在 sim.step 前覆盖 PhysX force。 + effort_target.mul_(self.motor_strength) + self.robot.data.applied_torque[:] = effort_target + self.robot.root_physx_view.set_dof_actuation_forces(effort_target, all_indices) def get_observations(self): actor_obs, critic_obs = self.compute_observations() From bcfb1f45eab2969660ba11749efc8e5d3594dfe0 Mon Sep 17 00:00:00 2001 From: m0_51228868 Date: Thu, 11 Jun 2026 09:36:29 +0800 Subject: [PATCH 6/6] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E8=AE=AD=E7=BB=83?= =?UTF-8?q?=E5=8F=82=E6=95=B0,=E5=AE=9E=E7=8E=B0=E6=96=B0=E7=9A=84?= =?UTF-8?q?=E8=AE=AD=E7=BB=83.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- WMP_NETWORK_DESIGN_DIFF_REPORT.txt | 76 +++++++++++++ legged_lab/envs/a1/a1_config.py | 6 +- legged_lab/envs/b2/b2_config.py | 64 +++++------ legged_lab/envs/base/base_env_config.py | 4 +- legged_lab/envs/rb160w/rb160w_config.py | 2 +- legged_lab/mdp/rewards.py | 28 +++++ legged_lab/runners/wmp_amp_runner.py | 30 ++++- legged_lab/scripts/inspect_rgbd_terrain.py | 28 ++--- legged_lab/scripts/inspect_wmp_edge_mask.py | 8 +- legged_lab/scripts/inspect_wmp_world_model.py | 16 +-- legged_lab/scripts/play.py | 32 +++--- legged_lab/scripts/train.py | 39 ++++++- legged_lab/sensors/__init__.py | 14 ++- .../sensors/wmp_partial_tiled_camera.py | 39 +++++-- legged_lab/utils/cli_args.py | 12 ++ legged_lab/world_models/wmp/preprocess.py | 4 +- scripts/train_a1_wmp_amp.sh | 104 ++++++++++++++++++ scripts/train_a1_wmp_amp_warm_start.sh | 25 +++++ 18 files changed, 430 insertions(+), 101 deletions(-) create mode 100755 scripts/train_a1_wmp_amp.sh create mode 100755 scripts/train_a1_wmp_amp_warm_start.sh diff --git a/WMP_NETWORK_DESIGN_DIFF_REPORT.txt b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt index 1a46c9e..9257db8 100644 --- a/WMP_NETWORK_DESIGN_DIFF_REPORT.txt +++ b/WMP_NETWORK_DESIGN_DIFF_REPORT.txt @@ -952,3 +952,79 @@ critic 维度已经对齐，下一步不是改维度，而是验收字段语义: obs normalizer、rsl_rl v5 storage/update 细节。 4. 记录 reset cause 分布、episode length 分布、AMP policy/expert score 分布。 5. 如果 trunk contact reset 占比异常，再评估是否为 A1 WMP 增加 current-force 接触终止路径。 + + +12.11 2026-06-09 代码级对齐复查结论 +----------------------------------- +复查范围: + - 原版: + /tmp/bytedance_WMP/rsl_rl/runners/wmp_runner.py + /tmp/bytedance_WMP/rsl_rl/modules/actor_critic_wmp.py + /tmp/bytedance_WMP/rsl_rl/modules/depth_predictor.py + /tmp/bytedance_WMP/rsl_rl/algorithms/amp_ppo.py + /tmp/bytedance_WMP/rsl_rl/algorithms/amp_discriminator.py + /tmp/bytedance_WMP/dreamer/configs.yaml + /tmp/bytedance_WMP/legged_gym/envs/a1/a1_amp_config.py + /tmp/bytedance_WMP/legged_gym/envs/base/legged_robot.py + - 当前: + legged_lab/models/wmp_mlp_model.py + legged_lab/world_models/wmp/{config.py,models.py,controller.py,depth_predictor.py} + legged_lab/runners/wmp_amp_runner.py + legged_lab/algorithms/wmp_amp_ppo.py + legged_lab/amp/discriminator.py + legged_lab/envs/a1/a1_config.py + legged_lab/envs/base/base_env.py + legged_lab/mdp/events.py + legged_lab/sensors/wmp_partial_tiled_camera.py + +本轮结论: + - 未发现新的必须修改代码的 shape / 公式级不一致。 + - Actor/Critic/DepthPredictor/RSSM 的主要网络 shape 与原版配置继续保持一致。 + - WMPFixedEpisodeReplayBuffer 当前实现已按原版主要采样语义工作: + p_i = dataset_size_i / sum_j(dataset_size_j) + batch_length = min(sampled_dataset_size.min(), config.batch_length) + end_idx 从 [batch_length, dataset_size_i] 采样 + out["is_first"][:, 0] = 1，其余帧为 0。 + - controller.after_env_step() 对 reset 和 WMP tick 同时发生的处理顺序，与原版 runner 对齐: + 先把已结束 env 的旧 wm_buffer 写入 dataset； + 再把 wm_is_first/reset env 置 1； + 当前 tick 只写 not_reset_env_ids = (1 - wm_is_first)。 + 因此 reset 后的新初始状态不会被写入旧 episode。 + - DepthPredictor 结构与原版一致；当前 forward 额外支持 leading shape reshape， + 这是 API 兼容增强，不改变单 batch / batch-time 的输出语义。 + - AMPDiscriminator 的 reward 公式、task_reward_lerp 实际行为、grad penalty 公式与原版一致。 + 原版虽然叫 lerp，但实现为 disc_r + task_r；当前也保留这一行为。 + - A1 WMP PPO/AMP 关键配置仍与原版一致: + entropy_coef=0.01 + vel_predict_coef=1.0 + num_learning_epochs=5 + num_mini_batches=4 + amp_reward_coef=0.01 + amp_task_reward_lerp=0.3 + amp_discr_hidden_dims=[1024,512] + min_normalized_std=[0.05,0.02,0.05]*4 + +确认保留、后续 review 的 API/框架差异: + - WMPAMPRunner 基于 rsl_rl v5 actor/critic/storage 接口，不是原版 rsl_rl runner 逐行迁移。 + - checkpoint key 更完整但不兼容原版；继续保留。 + - ContactSensor、相机、材质、质量、COM、显式 actuator torque buffer 都来自 IsaacLab/PhysX 路径， + 不强行改回 IsaacGym 逐行实现。 + - _refresh_wmp_privileged_buffers_from_sim() 会从 IsaacLab sim 状态同步 privileged buffer。 + 这属于 IsaacLab 后端兜底同步；主要语义仍以 wmp_recording_* EventTerm 的同源采样为准。 + - action_delay 的插入位置仍是 sim_step 级 position target 选择；原版是在 _compute_torques() 入口选择 + last_actions/actions。该项继续作为 API/框架差异保留待 review。 + - AMPReplayBuffer/Normalizer 使用 torch 实现；行为与原版 numpy/cpu 路径相近，但不是逐行相同。 + +本轮没有改代码的原因: + - 已检查到的差异均属于报告中已标注的 IsaacLab/rsl_rl v5/API 迁移差异、日志命名差异或待统计验证项。 + - 没有发现类似 12.5 中 __init__ 参数缺失、采样边界错误、维度错误这类硬 bug。 + +验证情况: + - 本轮尝试用: + source /home/tower/miniconda/etc/profile.d/conda.sh + conda activate isaaclab + /home/tower/Bags/IsaacLab/isaaclab.sh -p - + 做轻量配置 import 检查，但当前 shell 环境缺少 pxr 模块，报: + ModuleNotFoundError: No module named 'pxr' + - 因此本轮只完成静态源码对齐复查，未运行 IsaacLab runtime shape 检查。 + - 下一次 runtime 验收仍建议沿用报告开头的 headless 检查脚本，并补充 replay/depth 统计打印。 diff --git a/legged_lab/envs/a1/a1_config.py b/legged_lab/envs/a1/a1_config.py index c6b7814..86128e6 100644 --- a/legged_lab/envs/a1/a1_config.py +++ b/legged_lab/envs/a1/a1_config.py @@ -381,9 +381,9 @@ def __post_init__(self): self.policy.actor_hidden_dims = [256, 128, 64] self.policy.critic_hidden_dims = [512, 256, 128] self.algorithm.class_name = "legged_lab.algorithms.wmp_amp_ppo:WMPAMPPPO" - self.algorithm.learning_rate = 3.0e-4 - self.algorithm.schedule = "fixed" - self.algorithm.desired_kl = 0.006 + self.algorithm.learning_rate = 1.0e-3 + self.algorithm.schedule = "adaptive" + self.algorithm.desired_kl = 0.01 self.algorithm.entropy_coef = 0.01 self.algorithm.num_learning_epochs = 5 self.algorithm.num_mini_batches = 4 diff --git a/legged_lab/envs/b2/b2_config.py b/legged_lab/envs/b2/b2_config.py index e0c6b8a..152c970 100644 --- a/legged_lab/envs/b2/b2_config.py +++ b/legged_lab/envs/b2/b2_config.py @@ -419,19 +419,19 @@ def __post_init__(self): super().__post_init__() self.scene.robot = B2_CFG self.scene.height_scanner.prim_body_name = "base_link" - self.scene.gemini2_camera.enable = True - self.scene.gemini2_camera.spawn_prim_path = "base_link/wmp_depth_camera" - self.scene.gemini2_camera.model_name = "wmp_front_depth" - self.scene.gemini2_camera.camera_model = "pinhole" - self.scene.gemini2_camera.spawn_offset_pos = (0.27, 0.0, 0.03) - self.scene.gemini2_camera.spawn_offset_rot = (1.0, 0.0, 0.0, 0.0) - self.scene.gemini2_camera.width = 64 - self.scene.gemini2_camera.height = 64 - self.scene.gemini2_camera.depth_near = 0.0 - self.scene.gemini2_camera.depth_far = 2.0 - self.scene.gemini2_camera.horizontal_fov_deg = 58.0 - self.scene.gemini2_camera.randomize_rotation = True - self.scene.gemini2_camera.random_pitch_deg = (-5.0, 5.0) + self.scene.rgbd_camera.enable = True + self.scene.rgbd_camera.spawn_prim_path = "base_link/wmp_depth_camera" + self.scene.rgbd_camera.model_name = "wmp_front_depth" + self.scene.rgbd_camera.camera_model = "pinhole" + self.scene.rgbd_camera.spawn_offset_pos = (0.27, 0.0, 0.03) + self.scene.rgbd_camera.spawn_offset_rot = (1.0, 0.0, 0.0, 0.0) + self.scene.rgbd_camera.width = 64 + self.scene.rgbd_camera.height = 64 + self.scene.rgbd_camera.depth_near = 0.0 + self.scene.rgbd_camera.depth_far = 2.0 + self.scene.rgbd_camera.horizontal_fov_deg = 58.0 + self.scene.rgbd_camera.randomize_rotation = True + self.scene.rgbd_camera.random_pitch_deg = (-5.0, 5.0) @configclass @@ -453,10 +453,10 @@ def __post_init__(self): self.sim.physx.gpu_max_rigid_patch_count = 2 * 2**15 self.scene.terrain_type = "plane" self.scene.terrain_generator = None - self.scene.gemini2_camera.enable = False - self.scene.gemini2_camera.enable_rgb = False - self.scene.gemini2_camera.enable_depth = False - self.scene.gemini2_camera.allow_missing_depth_fallback = True + self.scene.rgbd_camera.enable = False + self.scene.rgbd_camera.enable_rgb = False + self.scene.rgbd_camera.enable_depth = False + self.scene.rgbd_camera.allow_missing_depth_fallback = True self.normalization.obs_scales.lin_vel = 1.0 self.normalization.obs_scales.ang_vel = 0.25 self.normalization.obs_scales.projected_gravity = 1.0 @@ -529,10 +529,10 @@ def __post_init__(self): self.sim.physx.gpu_max_rigid_patch_count = 2 * 2**15 self.scene.terrain_type = "plane" self.scene.terrain_generator = None - self.scene.gemini2_camera.enable = False - self.scene.gemini2_camera.enable_rgb = False - self.scene.gemini2_camera.enable_depth = False - self.scene.gemini2_camera.allow_missing_depth_fallback = True + self.scene.rgbd_camera.enable = False + self.scene.rgbd_camera.enable_rgb = False + self.scene.rgbd_camera.enable_depth = False + self.scene.rgbd_camera.allow_missing_depth_fallback = True self.normalization.obs_scales.lin_vel = 1.0 self.normalization.obs_scales.ang_vel = 0.25 self.normalization.obs_scales.projected_gravity = 1.0 @@ -599,7 +599,7 @@ def __post_init__(self): super().__post_init__() self.scene.robot = B2_CFG self.scene.height_scanner.prim_body_name = "base_link" - self.scene.gemini2_camera.enable = True + self.scene.rgbd_camera.enable = True @configclass @@ -617,16 +617,16 @@ class B2RGBDWMPAMPFlatEnvCfg(B2RGBDFlatEnvCfg): def __post_init__(self): super().__post_init__() wmp_update_interval = 5 - self.scene.gemini2_camera.enable = True - self.scene.gemini2_camera.enable_rgb = False - self.scene.gemini2_camera.enable_depth = True - self.scene.gemini2_camera.partial_camera = True - self.scene.gemini2_camera.partial_camera_num_envs = 1024 - self.scene.gemini2_camera.partial_camera_seed = 42 - self.scene.gemini2_camera.partial_camera_force_tilt_crawl = True - self.scene.gemini2_camera.model_name = "wmp_front_depth" - self.scene.gemini2_camera.update_period = self.sim.dt * self.sim.decimation * wmp_update_interval - self.scene.gemini2_camera.allow_missing_depth_fallback = True + self.scene.rgbd_camera.enable = True + self.scene.rgbd_camera.enable_rgb = False + self.scene.rgbd_camera.enable_depth = True + self.scene.rgbd_camera.partial_camera = True + self.scene.rgbd_camera.partial_camera_num_envs = 1024 + self.scene.rgbd_camera.partial_camera_seed = 42 + self.scene.rgbd_camera.partial_camera_force_tilt_crawl = True + self.scene.rgbd_camera.model_name = "wmp_front_depth" + self.scene.rgbd_camera.update_period = self.sim.dt * self.sim.decimation * wmp_update_interval + self.scene.rgbd_camera.allow_missing_depth_fallback = True self.sim.render_interval = self.sim.decimation * wmp_update_interval diff --git a/legged_lab/envs/base/base_env_config.py b/legged_lab/envs/base/base_env_config.py index 8e7367e..875dd78 100644 --- a/legged_lab/envs/base/base_env_config.py +++ b/legged_lab/envs/base/base_env_config.py @@ -32,7 +32,6 @@ CommandsCfg, DomainRandCfg, EventCfg, - Gemini2CameraCfg, HeightScannerCfg, MotorStrengthCfg, NoiseCfg, @@ -42,6 +41,7 @@ PhysxCfg, RewardCfg, RewardSettingsCfg, + RGBDCameraCfg, RobotCfg, SimCfg, ) @@ -66,7 +66,7 @@ class BaseEnvCfg: debug_vis=False, drift_range=(0.0, 0.0), # (0.3, 0.3) ), - gemini2_camera=Gemini2CameraCfg(enable=False), + rgbd_camera=RGBDCameraCfg(enable=False), ) robot: RobotCfg = RobotCfg( actor_obs_history_length=10, diff --git a/legged_lab/envs/rb160w/rb160w_config.py b/legged_lab/envs/rb160w/rb160w_config.py index 91ba80c..190b14a 100644 --- a/legged_lab/envs/rb160w/rb160w_config.py +++ b/legged_lab/envs/rb160w/rb160w_config.py @@ -137,7 +137,7 @@ def __post_init__(self): self.scene.terrain_generator = None self.scene.height_scanner.enable_height_scan = False self.scene.height_scanner.prim_body_name = "base_link" - self.scene.gemini2_camera.enable = False + self.scene.rgbd_camera.enable = False self.scene.env_spacing = 3.0 self.robot.terminate_contacts_body_names = [".*base_link.*"] diff --git a/legged_lab/mdp/rewards.py b/legged_lab/mdp/rewards.py index c0ce4d6..e9017d2 100644 --- a/legged_lab/mdp/rewards.py +++ b/legged_lab/mdp/rewards.py @@ -421,6 +421,34 @@ def feet_stumble(env: BaseEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor: ) +def wmp_feet_stumble(env: BaseEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor: + """WMP 原版 `feet_stumble` 奖励。 + + 原版公式: + rew = any(||F_xy||_2 > 4 * |F_z|) + rew *= terrain_levels > 3 + rew[非 gap/pit 段地形] = 0 + + 其中 `gap_start_col` 和 `climb_end_col` 对应原版按 env 连续切片的 + `[gap_start_idx, pit_end_idx)` 地形段。当前用 terrain_types 的列号 gate + 表达同一语义。 + """ + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + forces = contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids] + rew = torch.any(torch.norm(forces[..., :2], dim=-1) > 4.0 * torch.abs(forces[..., 2]), dim=1) + + terrain_levels = getattr(env, "terrain_levels", None) + terrain_types = getattr(env, "terrain_types", None) + if terrain_levels is not None: + rew &= terrain_levels > 3 + if terrain_types is not None: + type_gate = (terrain_types >= getattr(env, "gap_start_col", 0)) & ( + terrain_types < getattr(env, "climb_end_col", terrain_types.max().item() + 1) + ) + rew &= type_gate + return rew.float() + + def feet_too_near_humanoid( env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), threshold: float = 0.2 ) -> torch.Tensor: diff --git a/legged_lab/runners/wmp_amp_runner.py b/legged_lab/runners/wmp_amp_runner.py index 5763e47..4d1fa24 100644 --- a/legged_lab/runners/wmp_amp_runner.py +++ b/legged_lab/runners/wmp_amp_runner.py @@ -6,6 +6,7 @@ import glob import os import resource +import statistics import time import torch @@ -19,6 +20,11 @@ from legged_lab.amp import AMPDiscriminator, AMPLoader, Normalizer from legged_lab.world_models.wmp import WMPTrainingController, WorldModel, make_default_wmp_config +try: + import wandb +except ImportError: # pragma: no cover - wandb is optional. + wandb = None + class WMPAMPRunner: def __init__(self, env, train_cfg: dict, log_dir: str | None = None, device: str = "cpu"): @@ -220,7 +226,9 @@ def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = Fals for it in range(start_it, total_it): curriculum_logs = {} - if hasattr(self.env, "update_reward_curriculum"): + if hasattr(self.env, "update_training_curriculum"): + curriculum_logs = self.env.update_training_curriculum(it) + elif hasattr(self.env, "update_reward_curriculum"): curriculum_logs = self.env.update_reward_curriculum(it) start = time.time() with torch.inference_mode(): @@ -284,6 +292,7 @@ def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = Fals self.alg.get_policy().output_std, None, ) + self._log_wandb_history(it, collect_time, learn_time, loss_dict) if self.logger.writer is not None and it % self.cfg["save_interval"] == 0: self.save(os.path.join(self.logger.log_dir, f"model_{it}.pt")) if self.logger.writer is not None: @@ -308,6 +317,25 @@ def _resource_stats(self) -> dict[str, float]: stats["gpu_reserved_gb"] = torch.cuda.memory_reserved(device) / (1024.0**3) return stats + def _log_wandb_history(self, it: int, collect_time: float, learn_time: float, loss_dict: dict): + if self.cfg.get("logger") != "wandb" or wandb is None or wandb.run is None: + return + collection_size = self.cfg["num_steps_per_env"] * self.env.num_envs * self.gpu_world_size + payload = { + "Perf/total_fps_batched": int(collection_size / max(collect_time + learn_time, 1.0e-8)), + "Perf/collection_time_batched": collect_time, + "Perf/learning_time_batched": learn_time, + "Policy/mean_std_batched": self.alg.get_policy().output_std.mean().item(), + } + for key, value in loss_dict.items(): + if isinstance(value, torch.Tensor): + value = value.detach().float().mean().item() + payload[f"Loss/{key}_batched"] = float(value) + if len(getattr(self.logger, "rewbuffer", [])) > 0: + payload["Train/mean_reward_batched"] = statistics.mean(self.logger.rewbuffer) + payload["Train/mean_episode_length_batched"] = statistics.mean(self.logger.lenbuffer) + wandb.log(payload, step=it, commit=True) + def load(self, path: str, load_cfg: dict | None = None, strict: bool = True, map_location: str | None = None): loaded = torch.load(path, weights_only=False, map_location=map_location) if self.alg.load(loaded, load_cfg, strict): diff --git a/legged_lab/scripts/inspect_rgbd_terrain.py b/legged_lab/scripts/inspect_rgbd_terrain.py index 4721677..a2853fc 100644 --- a/legged_lab/scripts/inspect_rgbd_terrain.py +++ b/legged_lab/scripts/inspect_rgbd_terrain.py @@ -68,27 +68,27 @@ def main(): if args_cli.device is not None: env_cfg.device = args_cli.device - env_cfg.scene.gemini2_camera.enable = True - env_cfg.scene.gemini2_camera.enable_rgb = True - env_cfg.scene.gemini2_camera.enable_depth = True - env_cfg.scene.gemini2_camera.partial_camera = False + env_cfg.scene.rgbd_camera.enable = True + env_cfg.scene.rgbd_camera.enable_rgb = True + env_cfg.scene.rgbd_camera.enable_depth = True + env_cfg.scene.rgbd_camera.partial_camera = False if args_cli.camera_model is not None: - env_cfg.scene.gemini2_camera.camera_model = args_cli.camera_model + env_cfg.scene.rgbd_camera.camera_model = args_cli.camera_model if args_cli.camera_name is not None: - env_cfg.scene.gemini2_camera.model_name = args_cli.camera_name + env_cfg.scene.rgbd_camera.model_name = args_cli.camera_name if args_cli.camera_prim_path is not None: - env_cfg.scene.gemini2_camera.spawn_prim_path = args_cli.camera_prim_path + env_cfg.scene.rgbd_camera.spawn_prim_path = args_cli.camera_prim_path if args_cli.camera_offset_pos is not None: - env_cfg.scene.gemini2_camera.spawn_offset_pos = tuple(args_cli.camera_offset_pos) + env_cfg.scene.rgbd_camera.spawn_offset_pos = tuple(args_cli.camera_offset_pos) if args_cli.camera_offset_rot is not None: - env_cfg.scene.gemini2_camera.spawn_offset_rot = tuple(args_cli.camera_offset_rot) - camera_prim_path = "{ENV_REGEX_NS}/Robot/" + env_cfg.scene.gemini2_camera.spawn_prim_path.strip("/") + env_cfg.scene.rgbd_camera.spawn_offset_rot = tuple(args_cli.camera_offset_rot) + camera_prim_path = "{ENV_REGEX_NS}/Robot/" + env_cfg.scene.rgbd_camera.spawn_prim_path.strip("/") env_class = task_registry.get_task_class(args_cli.task) env = env_class(env_cfg, args_cli.headless) env.sim.set_camera_view(eye=[2.8, -2.8, 1.6], target=[0.0, 0.0, 0.45]) - depth_camera = env.scene.sensors["gemini2_depth_camera"] - rgb_camera = env.scene.sensors.get("gemini2_rgb_camera", depth_camera) + depth_camera = env.scene.sensors["rgbd_camera"] + rgb_camera = env.scene.sensors.get("rgbd_camera", depth_camera) print(f"[INFO] reading spawned RGBD camera prim: {camera_prim_path}") rgb = None @@ -119,8 +119,8 @@ def main(): save_images_to_file( _normalize_depth( depth, - env_cfg.scene.gemini2_camera.depth_near, - env_cfg.scene.gemini2_camera.depth_far, + env_cfg.scene.rgbd_camera.depth_near, + env_cfg.scene.rgbd_camera.depth_far, ), str(depth_path), ) diff --git a/legged_lab/scripts/inspect_wmp_edge_mask.py b/legged_lab/scripts/inspect_wmp_edge_mask.py index 3902f89..b3b156e 100644 --- a/legged_lab/scripts/inspect_wmp_edge_mask.py +++ b/legged_lab/scripts/inspect_wmp_edge_mask.py @@ -29,10 +29,10 @@ def main(): env_cfg.scene.terrain_type = "generator" env_cfg.scene.terrain_generator = WMP_TERRAIN_CFGS[args_cli.wmp_terrain] env_cfg.scene.max_init_terrain_level = 0 - env_cfg.scene.gemini2_camera.enable = False - env_cfg.scene.gemini2_camera.enable_rgb = False - env_cfg.scene.gemini2_camera.enable_depth = False - env_cfg.scene.gemini2_camera.allow_missing_depth_fallback = True + env_cfg.scene.rgbd_camera.enable = False + env_cfg.scene.rgbd_camera.enable_rgb = False + env_cfg.scene.rgbd_camera.enable_depth = False + env_cfg.scene.rgbd_camera.allow_missing_depth_fallback = True env_cfg.noise.add_noise = False if args_cli.device is not None: diff --git a/legged_lab/scripts/inspect_wmp_world_model.py b/legged_lab/scripts/inspect_wmp_world_model.py index baff383..f01b589 100644 --- a/legged_lab/scripts/inspect_wmp_world_model.py +++ b/legged_lab/scripts/inspect_wmp_world_model.py @@ -10,7 +10,7 @@ from legged_lab.utils import task_registry -parser = argparse.ArgumentParser(description="Inspect WMP RSSM world model with b2_rgbd Gemini2 depth.") +parser = argparse.ArgumentParser(description="Inspect WMP RSSM world model with b2_rgbd RGBD depth.") parser.add_argument("--task", type=str, default="b2_rgbd_rough", help="Task name.") parser.add_argument("--num_envs", type=int, default=1, help="Number of environments.") parser.add_argument("--steps", type=int, default=5, help="Number of zero-action warmup steps.") @@ -40,7 +40,7 @@ def main(): env_cfg.commands.ranges.lin_vel_y = (0.0, 0.0) env_cfg.commands.ranges.ang_vel_z = (0.0, 0.0) env_cfg.commands.ranges.heading = (0.0, 0.0) - env_cfg.scene.gemini2_camera.enable = True + env_cfg.scene.rgbd_camera.enable = True if args_cli.device is not None: env_cfg.device = args_cli.device @@ -50,7 +50,7 @@ def main(): exit_code = 1 try: env.sim.set_camera_view(eye=[2.8, -2.8, 1.6], target=[0.0, 0.0, 0.45]) - depth_camera = env.scene.sensors["gemini2_depth_camera"] + depth_camera = env.scene.sensors["rgbd_camera"] depth = None for step_idx in range(args_cli.steps): @@ -61,19 +61,19 @@ def main(): assert depth is not None wm_device = args_cli.wm_device or env.device - _log(f"preprocessing Gemini2 depth on device={wm_device}") + _log(f"preprocessing RGBD depth on device={wm_device}") prop = env.get_wmp_proprioception().to(wm_device) forward_height_map = env.get_wmp_forward_height_map().to(wm_device) camera_env_ids = env.get_depth_camera_env_ids() depth_nchw = depth_to_nchw( depth.to(wm_device), - near=env_cfg.scene.gemini2_camera.depth_near, - far=env_cfg.scene.gemini2_camera.depth_far, + near=env_cfg.scene.rgbd_camera.depth_near, + far=env_cfg.scene.rgbd_camera.depth_far, ) camera_image = depth_to_wmp_image( depth.to(wm_device), - near=env_cfg.scene.gemini2_camera.depth_near, - far=env_cfg.scene.gemini2_camera.depth_far, + near=env_cfg.scene.rgbd_camera.depth_near, + far=env_cfg.scene.rgbd_camera.depth_far, ) if camera_image.shape[0] == env.num_envs: image = camera_image diff --git a/legged_lab/scripts/play.py b/legged_lab/scripts/play.py index 59e2ffb..419bc02 100644 --- a/legged_lab/scripts/play.py +++ b/legged_lab/scripts/play.py @@ -44,7 +44,7 @@ ) parser.add_argument("--depth_image_dir", type=str, default=None, help="Directory for saved WMP depth image PNGs.") parser.add_argument("--depth_image_save_interval", type=int, default=10, help="Save one depth image every N play steps.") -parser.add_argument("--show_depth_points", action="store_true", help="Visualize Gemini2 depth hits as red debug points.") +parser.add_argument("--show_depth_points", action="store_true", help="Visualize RGBD depth hits as red debug points.") parser.add_argument("--show_height_scan_points", action="store_true", help="Visualize height scanner ray hits.") parser.add_argument("--enable_play_push", action="store_true", help="Keep interval push disturbances enabled during play.") parser.add_argument( @@ -405,20 +405,20 @@ def play(): if args_cli.num_envs is not None: env_cfg.scene.num_envs = args_cli.num_envs if args_cli.show_depth_points or args_cli.show_depth_image: - env_cfg.scene.gemini2_camera.enable = True - env_cfg.scene.gemini2_camera.enable_depth = True + env_cfg.scene.rgbd_camera.enable = True + env_cfg.scene.rgbd_camera.enable_depth = True if args_cli.camera_offset_pos is not None: - env_cfg.scene.gemini2_camera.spawn_offset_pos = tuple(args_cli.camera_offset_pos) + env_cfg.scene.rgbd_camera.spawn_offset_pos = tuple(args_cli.camera_offset_pos) if args_cli.camera_offset_rot is not None: - env_cfg.scene.gemini2_camera.spawn_offset_rot = tuple(args_cli.camera_offset_rot) + env_cfg.scene.rgbd_camera.spawn_offset_rot = tuple(args_cli.camera_offset_rot) if args_cli.camera_random_pitch_deg is not None: - env_cfg.scene.gemini2_camera.random_pitch_deg = tuple(args_cli.camera_random_pitch_deg) + env_cfg.scene.rgbd_camera.random_pitch_deg = tuple(args_cli.camera_random_pitch_deg) if args_cli.camera_fov_deg is not None: - env_cfg.scene.gemini2_camera.horizontal_fov_deg = float(args_cli.camera_fov_deg) + env_cfg.scene.rgbd_camera.horizontal_fov_deg = float(args_cli.camera_fov_deg) if args_cli.camera_disable_random_rotation: - env_cfg.scene.gemini2_camera.randomize_rotation = False - if env_cfg.scene.gemini2_camera.partial_camera: - env_cfg.scene.gemini2_camera.partial_camera_num_envs = env_cfg.scene.num_envs + env_cfg.scene.rgbd_camera.randomize_rotation = False + if env_cfg.scene.rgbd_camera.partial_camera: + env_cfg.scene.rgbd_camera.partial_camera_num_envs = env_cfg.scene.num_envs agent_cfg = update_rsl_rl_cfg(agent_cfg, args_cli) env_cfg.scene.seed = agent_cfg.seed @@ -500,10 +500,10 @@ def play(): keyboard = Keyboard(env) # noqa:F841 obs = env.get_observations() - depth_camera = env.scene.sensors.get("gemini2_depth_camera") if args_cli.show_depth_points or args_cli.show_depth_image else None + depth_camera = env.scene.sensors.get("rgbd_camera") if args_cli.show_depth_points or args_cli.show_depth_image else None depth_draw = _acquire_debug_draw_interface() if args_cli.show_depth_points else None if depth_camera is not None: - # TiledCamera 默认返回初始化时的相机位姿；红点需要跟随机器人上的 Gemini2， + # TiledCamera 默认返回初始化时的相机位姿；红点需要跟随机器人上的 RGBD， # 因此播放调试时开启最新位姿更新。 depth_camera.cfg.update_latest_camera_pose = True camera_ids = getattr(depth_camera, "camera_env_ids", None) @@ -528,8 +528,8 @@ def play(): if depth_camera is not None and args_cli.show_depth_image: depth_window_available = _handle_wmp_depth_image( depth_camera, - near=env_cfg.scene.gemini2_camera.depth_near, - far=env_cfg.scene.gemini2_camera.depth_far, + near=env_cfg.scene.rgbd_camera.depth_near, + far=env_cfg.scene.rgbd_camera.depth_far, mode=args_cli.depth_image_mode, output_dir=depth_image_dir, step=depth_image_step, @@ -540,8 +540,8 @@ def play(): if depth_camera is not None and depth_draw is not None: depth_points, depth_origins = _depth_to_hit_points( depth_camera, - near=env_cfg.scene.gemini2_camera.depth_near, - far=env_cfg.scene.gemini2_camera.depth_far, + near=env_cfg.scene.rgbd_camera.depth_near, + far=env_cfg.scene.rgbd_camera.depth_far, stride=args_cli.depth_point_stride, max_points=args_cli.depth_point_max, forward_min=args_cli.depth_point_forward_min, diff --git a/legged_lab/scripts/train.py b/legged_lab/scripts/train.py index f7830bd..fca9062 100644 --- a/legged_lab/scripts/train.py +++ b/legged_lab/scripts/train.py @@ -77,6 +77,13 @@ def setup_wandb_env(logger_name: str): os.environ["WANDB_API_KEY"] = args_cli.wandb_api_key +def _scene_rgbd_camera_cfg(env_cfg): + """读取当前统一的 RGBD camera 配置。""" + if hasattr(env_cfg.scene, "rgbd_camera"): + return env_cfg.scene.rgbd_camera + return None + + def train(): runner: OnPolicyRunner env = None @@ -89,14 +96,18 @@ def train(): if args_cli.num_envs is not None: env_cfg.scene.num_envs = args_cli.num_envs + rgbd_camera_cfg = _scene_rgbd_camera_cfg(env_cfg) if getattr(args_cli, "wmp_camera_num_envs", None) is not None: - env_cfg.scene.gemini2_camera.partial_camera_num_envs = args_cli.wmp_camera_num_envs + if rgbd_camera_cfg is None: + raise RuntimeError("--wmp_camera_num_envs requires an RGBD camera config on env_cfg.scene.") + rgbd_camera_cfg.partial_camera_num_envs = args_cli.wmp_camera_num_envs if args_cli.runner == "wmp_amp" and not getattr(args_cli, "enable_cameras", False): print("[WARN] WMP-AMP runner without --enable_cameras: using zero-depth fallback for dry smoke only.") - env_cfg.scene.gemini2_camera.enable = False - env_cfg.scene.gemini2_camera.enable_depth = False - env_cfg.scene.gemini2_camera.enable_rgb = False - env_cfg.scene.gemini2_camera.allow_missing_depth_fallback = True + if rgbd_camera_cfg is not None: + rgbd_camera_cfg.enable = False + rgbd_camera_cfg.enable_depth = False + rgbd_camera_cfg.enable_rgb = False + rgbd_camera_cfg.allow_missing_depth_fallback = True agent_cfg = update_rsl_rl_cfg(agent_cfg, args_cli) env_cfg.scene.seed = agent_cfg.seed @@ -146,6 +157,24 @@ def train(): runner = runner_cls(env, cfg_dict, log_dir=log_dir, device=agent_cfg.device) remaining_iterations = agent_cfg.max_iterations + warm_start_checkpoint = getattr(args_cli, "warm_start_checkpoint", None) + if warm_start_checkpoint and agent_cfg.resume: + raise RuntimeError("--warm_start_checkpoint cannot be used together with --resume.") + + if warm_start_checkpoint: + warm_start_checkpoint = os.path.abspath(warm_start_checkpoint) + print(f"[INFO] Warm-starting model weights from: {warm_start_checkpoint}") + runner.load( + warm_start_checkpoint, + load_cfg={"optimizer": bool(getattr(args_cli, "warm_start_load_optimizer", False))}, + strict=True, + map_location=agent_cfg.device, + ) + runner.current_learning_iteration = 0 + print( + "[INFO] Warm-start loaded. Reset learning iteration to 0; " + f"optimizer_loaded={bool(getattr(args_cli, 'warm_start_load_optimizer', False))}." + ) if agent_cfg.resume: # get path to previous checkpoint diff --git a/legged_lab/sensors/__init__.py b/legged_lab/sensors/__init__.py index df34d45..9b85da1 100644 --- a/legged_lab/sensors/__init__.py +++ b/legged_lab/sensors/__init__.py @@ -1,6 +1,16 @@ # -*- coding: utf-8 -*- """LeggedLab 本地传感器扩展。""" -from .wmp_partial_tiled_camera import WMPPartialTiledCamera, WMPPartialTiledCameraCfg, select_wmp_camera_env_ids +from .wmp_partial_tiled_camera import ( + WMPPartialTiledCamera, + WMPPartialTiledCameraCfg, + select_wmp_camera_env_ids, + select_wmp_depth_indices, +) -__all__ = ["WMPPartialTiledCamera", "WMPPartialTiledCameraCfg", "select_wmp_camera_env_ids"] +__all__ = [ + "WMPPartialTiledCamera", + "WMPPartialTiledCameraCfg", + "select_wmp_camera_env_ids", + "select_wmp_depth_indices", +] diff --git a/legged_lab/sensors/wmp_partial_tiled_camera.py b/legged_lab/sensors/wmp_partial_tiled_camera.py index f590c1f..7582203 100644 --- a/legged_lab/sensors/wmp_partial_tiled_camera.py +++ b/legged_lab/sensors/wmp_partial_tiled_camera.py @@ -208,13 +208,11 @@ def __init__(self, cfg: WMPPartialTiledCameraCfg): if camera_ids.numel() == 0: raise ValueError("WMPPartialTiledCamera requires at least one camera env id.") ordered_num_envs = int(cfg.full_num_envs or int(camera_ids.max().item()) + 1) - self.camera_env_ids = _ordered_unique_valid(camera_ids, ordered_num_envs, torch.device("cpu")) + self.requested_camera_env_ids = _ordered_unique_valid(camera_ids, ordered_num_envs, torch.device("cpu")) + self.camera_env_ids = self.requested_camera_env_ids.clone() full_num_envs = int(cfg.full_num_envs or (int(self.camera_env_ids.max().item()) + 1)) self.full_num_envs = full_num_envs - self.camera_env_mask = torch.zeros(full_num_envs, dtype=torch.bool) - self.camera_env_mask[self.camera_env_ids] = True - self.camera_env_id_inverse = torch.full((full_num_envs,), -1, dtype=torch.long) - self.camera_env_id_inverse[self.camera_env_ids] = torch.arange(self.camera_env_ids.numel(), dtype=torch.long) + self._set_camera_env_ids(self.camera_env_ids) self._wmp_full_prim_path = cfg.prim_path self._wmp_subset_prim_path = self._camera_subset_prim_path(cfg.prim_path, self.camera_env_ids) spawn_cfg = cfg.spawn @@ -235,12 +233,31 @@ def _initialize_impl(self): if len(actual_env_ids) == self._view.count: actual = torch.tensor(actual_env_ids, dtype=torch.long) if actual.numel() != self.camera_env_ids.numel() or not torch.equal(actual, self.camera_env_ids): - requested_head = self.camera_env_ids[:16].tolist() - actual_head = actual[:16].tolist() - raise RuntimeError( - "WMPPartialTiledCamera image slot order must match depth_index exactly. " - f"requested_head={requested_head}, actual_head={actual_head}" - ) + requested_set = set(self.camera_env_ids.tolist()) + actual_set = set(actual.tolist()) + if actual.numel() == self.camera_env_ids.numel() and requested_set == actual_set: + print( + "[WARN] WMPPartialTiledCamera image slot order differs from requested depth_index; " + "using actual IsaacLab camera view order for depth_index mapping. " + f"requested_head={self.camera_env_ids[:16].tolist()}, actual_head={actual[:16].tolist()}" + ) + self._set_camera_env_ids(actual) + else: + missing = sorted(requested_set - actual_set)[:16] + extra = sorted(actual_set - requested_set)[:16] + raise RuntimeError( + "WMPPartialTiledCamera camera env set mismatch. " + f"requested_count={self.camera_env_ids.numel()}, actual_count={actual.numel()}, " + f"requested_head={self.camera_env_ids[:16].tolist()}, actual_head={actual[:16].tolist()}, " + f"missing_head={missing}, extra_head={extra}" + ) + + def _set_camera_env_ids(self, camera_env_ids: torch.Tensor): + self.camera_env_ids = torch.as_tensor(camera_env_ids, dtype=torch.long, device="cpu").flatten().clone() + self.camera_env_mask = torch.zeros(self.full_num_envs, dtype=torch.bool) + self.camera_env_mask[self.camera_env_ids] = True + self.camera_env_id_inverse = torch.full((self.full_num_envs,), -1, dtype=torch.long) + self.camera_env_id_inverse[self.camera_env_ids] = torch.arange(self.camera_env_ids.numel(), dtype=torch.long) def _spawn_partial_cameras_from_cfg(self, cfg: WMPPartialTiledCameraCfg, spawn_cfg): stage = get_current_stage() diff --git a/legged_lab/utils/cli_args.py b/legged_lab/utils/cli_args.py index 609aa81..028b27c 100644 --- a/legged_lab/utils/cli_args.py +++ b/legged_lab/utils/cli_args.py @@ -49,6 +49,18 @@ def add_rsl_rl_args(parser: argparse.ArgumentParser): arg_group.add_argument("--resume", type=bool, default=None, help="Whether to resume from a checkpoint.") arg_group.add_argument("--load_run", type=str, default=None, help="Name of the run folder to resume from.") arg_group.add_argument("--checkpoint", type=str, default=None, help="Checkpoint file to resume from.") + arg_group.add_argument( + "--warm_start_checkpoint", + type=str, + default=None, + help="Load model weights from this checkpoint but start a fresh training run from iteration 0.", + ) + arg_group.add_argument( + "--warm_start_load_optimizer", + action="store_true", + default=False, + help="Also load optimizer states during warm-start. Defaults to false.", + ) # -- logger arguments arg_group.add_argument( "--logger", type=str, default="wandb", choices={"wandb", "tensorboard", "neptune"}, help="Logger module to use." diff --git a/legged_lab/world_models/wmp/preprocess.py b/legged_lab/world_models/wmp/preprocess.py index daf46e8..9bc429c 100644 --- a/legged_lab/world_models/wmp/preprocess.py +++ b/legged_lab/world_models/wmp/preprocess.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -"""Gemini2 depth 到 WMP image 输入的预处理。""" +"""RGBD depth 到 WMP image 输入的预处理。""" import torch from torch.nn import functional as F @@ -39,5 +39,5 @@ def depth_to_wmp_image( far: float = 2.0, image_size: tuple[int, int] = (64, 64), ) -> torch.Tensor: - """将 Gemini2 depth 转换为 WMP 内部 NHWC 图像 B,64,64,1。""" + """将 RGBD depth 转换为 WMP 内部 NHWC 图像 B,64,64,1。""" return depth_to_nchw(depth, near=near, far=far, image_size=image_size).permute(0, 2, 3, 1).contiguous() diff --git a/scripts/train_a1_wmp_amp.sh b/scripts/train_a1_wmp_amp.sh new file mode 100755 index 0000000..667ba25 --- /dev/null +++ b/scripts/train_a1_wmp_amp.sh @@ -0,0 +1,104 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +ISAACLAB_DIR="${ISAACLAB_DIR:-/home/tower/Bags/IsaacLab}" +CONDA_SH="${CONDA_SH:-/home/tower/miniconda/etc/profile.d/conda.sh}" +CONDA_ENV="${CONDA_ENV:-isaaclab}" + +MODE="${1:-smoke}" +if [[ "$#" -gt 0 ]]; then + shift +fi + +case "${MODE}" in + smoke) + NUM_ENVS="${NUM_ENVS:-2}" + MAX_ITERATIONS="${MAX_ITERATIONS:-1}" + NUM_STEPS_PER_ENV="${NUM_STEPS_PER_ENV:-2}" + NUM_MINI_BATCHES="${NUM_MINI_BATCHES:-1}" + WMP_CAMERA_NUM_ENVS="${WMP_CAMERA_NUM_ENVS:-2}" + LOGGER="${LOGGER:-tensorboard}" + WANDB_MODE="${WANDB_MODE:-disabled}" + EXTRA_MODE_ARGS=( + --wmp_depth_training_iters "${WMP_DEPTH_TRAINING_ITERS:-1}" + --wmp_depth_batch_size "${WMP_DEPTH_BATCH_SIZE:-2}" + --wmp_train_steps_per_iter "${WMP_TRAIN_STEPS_PER_ITER:-1}" + --amp_num_preload_transitions "${AMP_NUM_PRELOAD_TRANSITIONS:-4096}" + ) + ;; + train) + NUM_ENVS="${NUM_ENVS:-4096}" + MAX_ITERATIONS="${MAX_ITERATIONS:-20000}" + NUM_STEPS_PER_ENV="${NUM_STEPS_PER_ENV:-24}" + NUM_MINI_BATCHES="${NUM_MINI_BATCHES:-4}" + WMP_CAMERA_NUM_ENVS="${WMP_CAMERA_NUM_ENVS:-1024}" + LOGGER="${LOGGER:-wandb}" + WANDB_MODE="${WANDB_MODE:-online}" + EXTRA_MODE_ARGS=() + ;; + dry) + NUM_ENVS="${NUM_ENVS:-2}" + MAX_ITERATIONS="${MAX_ITERATIONS:-1}" + NUM_STEPS_PER_ENV="${NUM_STEPS_PER_ENV:-2}" + NUM_MINI_BATCHES="${NUM_MINI_BATCHES:-1}" + WMP_CAMERA_NUM_ENVS="${WMP_CAMERA_NUM_ENVS:-0}" + LOGGER="${LOGGER:-tensorboard}" + WANDB_MODE="${WANDB_MODE:-disabled}" + EXTRA_MODE_ARGS=( + --wmp_depth_training_iters "${WMP_DEPTH_TRAINING_ITERS:-1}" + --wmp_depth_batch_size "${WMP_DEPTH_BATCH_SIZE:-2}" + --wmp_train_steps_per_iter "${WMP_TRAIN_STEPS_PER_ITER:-1}" + --amp_num_preload_transitions "${AMP_NUM_PRELOAD_TRANSITIONS:-4096}" + ) + ;; + *) + echo "Usage: $0 [smoke|train|dry] [extra train.py args...]" >&2 + exit 2 + ;; +esac + +if [[ -f "${CONDA_SH}" ]]; then + # shellcheck disable=SC1090 + source "${CONDA_SH}" + conda activate "${CONDA_ENV}" +fi + +if [[ -f "${ISAACLAB_DIR}/_isaac_sim/setup_conda_env.sh" ]]; then + # Isaac Sim 的运行时环境变量由这个脚本注入到当前 shell。 + # 公式无关；这是训练入口所需的环境准备。 + # shellcheck disable=SC1090 + set +u + source "${ISAACLAB_DIR}/_isaac_sim/setup_conda_env.sh" + set -u +fi + +cd "${ROOT_DIR}" +export PYTHONPATH="${ROOT_DIR}:${PYTHONPATH:-}" +export WANDB_MODE + +CMD=( + python legged_lab/scripts/train.py + --task a1_wmp_amp_terrain + --runner wmp_amp + --headless + --num_envs "${NUM_ENVS}" + --max_iterations "${MAX_ITERATIONS}" + --num_steps_per_env "${NUM_STEPS_PER_ENV}" + --num_mini_batches "${NUM_MINI_BATCHES}" + --logger "${LOGGER}" + --wandb_mode "${WANDB_MODE}" + --wmp_camera_num_envs "${WMP_CAMERA_NUM_ENVS}" + "${EXTRA_MODE_ARGS[@]}" +) + +if [[ "${MODE}" != "dry" ]]; then + CMD+=(--enable_cameras) +fi + +CMD+=("$@") + +echo "[INFO] mode=${MODE}" +echo "[INFO] root=${ROOT_DIR}" +echo "[INFO] command: ${CMD[*]}" +exec "${CMD[@]}" diff --git a/scripts/train_a1_wmp_amp_warm_start.sh b/scripts/train_a1_wmp_amp_warm_start.sh new file mode 100755 index 0000000..8ae8719 --- /dev/null +++ b/scripts/train_a1_wmp_amp_warm_start.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +DEFAULT_CKPT="${ROOT_DIR}/logs/a1_wmp_amp_terrain/2026-06-09_10-48-49/model_8000.pt" + +WARM_START_CHECKPOINT="${WARM_START_CHECKPOINT:-${1:-${DEFAULT_CKPT}}}" +if [[ "$#" -gt 0 ]]; then + shift +fi + +if [[ ! -f "${WARM_START_CHECKPOINT}" ]]; then + echo "[ERROR] Warm-start checkpoint not found: ${WARM_START_CHECKPOINT}" >&2 + exit 1 +fi + +export RUN_NAME="${RUN_NAME:-warm_start_from_model_8000_fixed_curriculum}" + +EXTRA_ARGS=(--warm_start_checkpoint "${WARM_START_CHECKPOINT}" --run_name "${RUN_NAME}") + +if [[ "${WARM_START_LOAD_OPTIMIZER:-0}" == "1" ]]; then + EXTRA_ARGS+=(--warm_start_load_optimizer) +fi + +exec "${ROOT_DIR}/scripts/train_a1_wmp_amp.sh" train "${EXTRA_ARGS[@]}" "$@"