Remove useless code

2026-04-07 11:06:02 +08:00
parent 090fcc093a
commit 7558467aca
13 changed files with 170 additions and 106 deletions
--- a/README.md
+++ b/README.md
@@ -28,7 +28,7 @@ Follow the step-by-step setup guide in [setup.md](doc/setup_en.md).
 Run the following command to launch training:

 ```bash
-python legged_gym/scripts/train.py --task=xxx
+python legged_gym/scripts/train.py --task=xxx --headless
 ```

 #### ⚙️  Arguments
@@ -59,14 +59,15 @@ The trained model above was evaluated using the [RoboGauge](https://github.com/w

 | Model | Score | Tracking  | Safety  | Quality  | Level | Download |
 | --- | --- | --- | --- | --- | --- | --- |
-| go2_moe_cts | **0.6819** | **0.6714** | **0.7794** | **0.7748** | **7.85** | [ckpt](https://drive.google.com/drive/folders/1aoXUxw-pGK1MbyzQ4IJzlA_tW8zrWP3Y?usp=drive_link) |
-| go2_moe_ng_cts | 0.6670 | 0.6552 | 0.7651 | 0.7613 | 7.67 | [ckpt](https://drive.google.com/drive/folders/1Rr89ZS0QJT-o-5LXsNqCWJdLGweqmN4Q?usp=drive_link) |
-| go2_ac_moe_cts | 0.6652 | 0.6527 | 0.7615 | 0.7552 | 7.57 | [ckpt](https://drive.google.com/file/d/1CDLsaR4XR3oG09ZHQ5u3lrJLfwyH2jz2/view?usp=drive_link) |
-| go2_mcp_cts | 0.6545 | 0.6440 | 0.7531 | 0.7476 | 7.48 | [ckpt](https://drive.google.com/drive/folders/1fd9cDVhV1dY6hcxuSZq2mcvFUp6V5Zfl?usp=drive_link) |
-| [HIM](https://github.com/InternRobotics/HIMLoco) | 0.5209 | 0.5200 | 0.6200 | 0.6100 | 5.78 | [ckpt](https://drive.google.com/file/d/1remJbGoTorqnArsz8Z1ewY4TVobss4Fb/view?usp=drive_link) |
-| [DreamWaQ](https://arxiv.org/abs/2301.10602) | 0.4832 | 0.4800 | 0.5800 | 0.5700 | 5.26 | [ckpt](https://drive.google.com/file/d/19BEBeiQqjHcPgGrN3AX6D7Yefs_8eswL/view?usp=drive_link) |
+| go2_moe_cts (Ours) | **0.6713** | **0.6669** | **0.7857** | **0.7392** | **7.85** | [ckpt](https://drive.google.com/drive/folders/1aoXUxw-pGK1MbyzQ4IJzlA_tW8zrWP3Y?usp=drive_link) |
+| go2_ac_moe_cts | 0.6509 | 0.6442 | 0.7644 | 0.7149 | 7.52 | [ckpt](https://drive.google.com/file/d/1CDLsaR4XR3oG09ZHQ5u3lrJLfwyH2jz2/view?usp=drive_link) |
+| go2_mcp_cts | 0.6399 | 0.6355 | 0.7542 | 0.7058 | 7.41 | [ckpt](https://drive.google.com/drive/folders/1fd9cDVhV1dY6hcxuSZq2mcvFUp6V5Zfl?usp=drive_link) |
+| go2_moe_ng_cts | 0.6519 | 0.6447 | 0.7639 | 0.7186 | 7.56 | [ckpt](https://drive.google.com/drive/folders/1Rr89ZS0QJT-o-5LXsNqCWJdLGweqmN4Q?usp=drive_link) |
+| [CTS](https://arxiv.org/pdf/2405.10830) vanilla | 0.5786 | 0.5755 | 0.7066 | 0.6624 | 6.83 | [ckpt]() |
+| [HIM](https://github.com/InternRobotics/HIMLoco) | 0.5379 | 0.5453 | 0.6476 | 0.6050 | 6.19 | [ckpt](https://drive.google.com/file/d/1remJbGoTorqnArsz8Z1ewY4TVobss4Fb/view?usp=drive_link) |
+| [DreamWaQ](https://arxiv.org/abs/2301.10602) | 0.5054 | 0.5105 | 0.6149 | 0.5730 | 5.74 | [ckpt](https://drive.google.com/file/d/19BEBeiQqjHcPgGrN3AX6D7Yefs_8eswL/view?usp=drive_link) |

-> In the downloaded ckpt, *.pt files are used for [Python deployment](#41-python-deployment), and *.onnx files are used for [C++ deployment](#42-c-deployment).
+> In the downloaded ckpt files, `*.pt` is used for [Python deployment](#41-python-deployment), and `*.onnx` is used for [C++ deployment](#42-c-deployment). The models above were all trained with self-collision disabled. In later tests, we found that enabling self-collision can also achieve strong results; see [go2_moe_cts_self_0.6669 - ckpt](https://drive.google.com/drive/folders/1znytqHNtDiZM5J4vaBd-EuM81l91D6s5?usp=drive_link).

 ### 2. Play

@@ -80,9 +81,9 @@ python legged_gym/scripts/play.py --task=xxx

 - Play launches on randomized terrain with difficulty between 7 and 9.
 - It automatically loads the latest checkpoint inside the experiment folder.
- Override via `experiment_name` and `checkpoint`, for example:
+- You can specify another model via `experiment_name`, `load_run`, and `checkpoint`, for example:
 	```bash
-	python legged_gym/scripts/play.py --task=go2_cts --num_envs 100 --experiment_name go2_cts_hard_terrain --checkpoint 100000
+	python legged_gym/scripts/play.py --task=go2_moe_cts --num_envs 100 --experiment_name go2_cts_hard_terrain --load_run Mar21_22-54-5-46_ --checkpoint 100000
 	```

 #### 💾 Policy Export
--- a/README_zh.md
+++ b/README_zh.md
@@ -28,7 +28,7 @@
 运行以下命令进行训练：

 ```bash
-python legged_gym/scripts/train.py --task=xxx
+python legged_gym/scripts/train.py --task=xxx --headless
 ```

 #### ⚙️  参数说明
@@ -59,14 +59,15 @@ python legged_gym/scripts/train.py --task=xxx

 | Model | Score | Tracking | Safety | Quality | Level | Download |
 | --- | --- | --- | --- | --- | --- | --- |
-| go2_moe_cts (Ours) | **0.6739** | **0.6647** | **0.7776** | **0.7739** | **7.85** | [ckpt](https://drive.google.com/drive/folders/1aoXUxw-pGK1MbyzQ4IJzlA_tW8zrWP3Y?usp=drive_link) |
-| go2_ac_moe_cts | 0.6541 | 0.6425 | 0.7558 | 0.7504 | 7.52 | [ckpt](https://drive.google.com/file/d/1CDLsaR4XR3oG09ZHQ5u3lrJLfwyH2jz2/view?usp=drive_link) |
-| go2_moe_ng_cts | 0.6537 | 0.6423 | 0.7554 | 0.7525 | 7.56 | [ckpt](https://drive.google.com/drive/folders/1Rr89ZS0QJT-o-5LXsNqCWJdLGweqmN4Q?usp=drive_link) |
-| go2_mcp_cts | 0.6423 | 0.6323 | 0.7464 | 0.7412 | 7.41 | [ckpt](https://drive.google.com/drive/folders/1fd9cDVhV1dY6hcxuSZq2mcvFUp6V5Zfl?usp=drive_link) |
-| [HIM](https://github.com/InternRobotics/HIMLoco) | 0.5401 | 0.5389 | 0.6412 | 0.6391 | 6.19 | [ckpt](https://drive.google.com/file/d/1remJbGoTorqnArsz8Z1ewY4TVobss4Fb/view?usp=drive_link) |
-| [DreamWaQ](https://arxiv.org/abs/2301.10602) | 0.5032 | 0.5010 | 0.6085 | 0.6032 | 5.74 | [ckpt](https://drive.google.com/file/d/19BEBeiQqjHcPgGrN3AX6D7Yefs_8eswL/view?usp=drive_link) |
+| go2_moe_cts (Ours) | **0.6713** | **0.6669** | **0.7857** | **0.7392** | **7.85** | [ckpt](https://drive.google.com/drive/folders/1aoXUxw-pGK1MbyzQ4IJzlA_tW8zrWP3Y?usp=drive_link) |
+| go2_ac_moe_cts | 0.6509 | 0.6442 | 0.7644 | 0.7149 | 7.52 | [ckpt](https://drive.google.com/file/d/1CDLsaR4XR3oG09ZHQ5u3lrJLfwyH2jz2/view?usp=drive_link) |
+| go2_mcp_cts | 0.6399 | 0.6355 | 0.7542 | 0.7058 | 7.41 | [ckpt](https://drive.google.com/drive/folders/1fd9cDVhV1dY6hcxuSZq2mcvFUp6V5Zfl?usp=drive_link) |
+| go2_moe_ng_cts | 0.6519 | 0.6447 | 0.7639 | 0.7186 | 7.56 | [ckpt](https://drive.google.com/drive/folders/1Rr89ZS0QJT-o-5LXsNqCWJdLGweqmN4Q?usp=drive_link) |
+| [CTS](https://arxiv.org/pdf/2405.10830) vanilla | 0.5786 | 0.5755 | 0.7066 | 0.6624 | 6.83 | [ckpt]() |
+| [HIM](https://github.com/InternRobotics/HIMLoco) | 0.5379 | 0.5453 | 0.6476 | 0.6050 | 6.19 | [ckpt](https://drive.google.com/file/d/1remJbGoTorqnArsz8Z1ewY4TVobss4Fb/view?usp=drive_link) |
+| [DreamWaQ](https://arxiv.org/abs/2301.10602) | 0.5054 | 0.5105 | 0.6149 | 0.5730 | 5.74 | [ckpt](https://drive.google.com/file/d/19BEBeiQqjHcPgGrN3AX6D7Yefs_8eswL/view?usp=drive_link) |

-> 下载的ckpt中*.pt用于[py部署](#41-python实物部署)，*.onnx用于[cpp部署](#42-c实物部署)
+> 下载的 ckpt 中，`*.pt` 用于[Python 实物部署](#41-python实物部署)，`*.onnx` 用于[C++ 实物部署](#42-c实物部署)。上述模型均在关闭自碰撞的设置下训练；后续测试发现，开启自碰撞也能取得不错效果，参考 [go2_moe_cts_self_0.6669 - ckpt](https://drive.google.com/drive/folders/1znytqHNtDiZM5J4vaBd-EuM81l91D6s5?usp=drive_link)。

 ---

@@ -82,9 +83,9 @@ python legged_gym/scripts/play.py --task=xxx

 - Play 启动参数为随机地形，难度在7到9之间。
 - 默认加载实验文件夹最新训练的一个模型。
- 可通过 `experiment_name` 和 `checkpoint` 指定其他模型，例如
+- 可通过 `experiment_name`, `load_run` 和 `checkpoint` 指定其他模型，例如
    ```bash
-    python legged_gym/scripts/play.py --task=go2_cts --num_envs 100 --experiment_name go2_cts_hard_terrain --checkpoint 100000
+    python legged_gym/scripts/play.py --task=go2_moe_cts --num_envs 100 --experiment_name go2_cts_hard_terrain --load_run Mar21_22-54-5-46_ --checkpoint 100000
    ```

 #### 💾 导出网络
@@ -196,4 +197,3 @@ python deploy_real_go2.py eth0
 新增内容根据 [MIT License](./LICENSE) 授权，原仓库unitree_rl_gym根据 [BSD 3-Clause License](./LICENSE) 授权。

 详情请阅读完整 [LICENSE 文件](./LICENSE)。
-
--- a/UPDATE.md
+++ b/UPDATE.md
@@ -1,3 +1,11 @@
+# 20260403
+## v1.0.3
+1. 修复last_last_action重置问题，修复resample_command在计算奖励前的问题
+2. 将所有配置中加入自碰撞
+# 20260325
+## v1.0.2-rc2
+1. 修复robogauge评估中返回None导致的训练中断问题
+2. 将自碰撞打开，真机表现更好
 # 20260126
 ## v1.0.2-rc1
 1. 修改高速移动的训练文件到最终版，删除配置中无用注释
--- a/legged_gym/envs/base/legged_robot.py
+++ b/legged_gym/envs/base/legged_robot.py
@@ -129,6 +129,9 @@ class LeggedRobot(BaseTask):
        # compute observations, rewards, resets, ...
        self.check_termination()
        self.compute_reward()
+        # resample commands must after reward computing
+        resampling_env_ids = ((self.commands_resampling_step <= 0.0) * (self.episode_length_buf < self.max_episode_length - 1)).nonzero(as_tuple=False).flatten()
+        self._resample_commands(resampling_env_ids)
        env_ids = self.reset_buf.nonzero(as_tuple=False).flatten()
        self.reset_idx(env_ids)
        
@@ -137,6 +140,7 @@ class LeggedRobot(BaseTask):

        self.compute_observations() # in some cases a simulation step might be required to refresh some obs (for example body positions)

+        self.last_last_actions[:] = self.last_actions[:]
        self.last_actions[:] = self.actions[:]
        self.last_dof_vel[:] = self.dof_vel[:]
        self.last_root_vel[:] = self.root_states[:, 7:13]
@@ -180,7 +184,7 @@ class LeggedRobot(BaseTask):
    def reset_idx(self, env_ids):
        """ Reset some environments.
            Calls self._reset_dofs(env_ids), self._reset_root_states(env_ids), and self._resample_commands(env_ids)
-            [Optional] calls self._update_terrain_curriculum(env_ids), self.update_command_curriculum(env_ids) and
+            [Optional] calls self._update_terrain_curriculum(env_ids),
            Logs episode info
            Resets some buffers

@@ -216,14 +220,13 @@ class LeggedRobot(BaseTask):
        # reset buffers
        self.actions[env_ids] = 0.
        self.last_actions[env_ids] = 0.
+        self.last_last_actions[env_ids] = 0.
        self.last_dof_vel[env_ids] = 0.
        self.feet_air_time[env_ids] = 0.
        self.episode_length_buf[env_ids] = 0
        self.reset_buf[env_ids] = 1
        self.commands_resampling_step[env_ids] = self.cfg.commands.resampling_time / self.dt
        self.commands_xy_accumulation[env_ids] = 0.0
-        if self.cfg.commands.curriculum:
-            self.update_command_curriculum(env_ids)
        self._resample_commands(env_ids)
        # fill extras
        self.extras["episode"] = {}
@@ -238,8 +241,7 @@ class LeggedRobot(BaseTask):
        for key in self.episode_sums.keys():
            self.extras["episode"]['rew_' + key] = torch.mean(self.episode_sums[key][env_ids]) / self.max_episode_length_s
            self.episode_sums[key][env_ids] = 0.
-        if self.cfg.commands.curriculum:
-            self.extras["episode"]["max_command_x"] = self.command_ranges["lin_vel_x"][1]
+        self.extras["episode"]["max_command_x"] = self.command_ranges["lin_vel_x"][1]
        # send timeout info to the algorithm
        if self.cfg.env.send_timeouts:
            self.extras["time_outs"] = self.time_out_buf
@@ -403,20 +405,8 @@ class LeggedRobot(BaseTask):
    
    def _post_physics_step_callback(self):
        """ Callback called before computing terminations, rewards, and observations
-            Default behaviour: Compute ang vel command based on target and heading, compute measured terrain heights and randomly push robots
+            Default behaviour: Compute measured terrain heights and randomly push robots
        """
-        # env_ids = (self.episode_length_buf % int(self.cfg.commands.resampling_time / self.dt)==0).nonzero(as_tuple=False).flatten()
-        resampling_env_ids = ((self.commands_resampling_step <= 0.0) * (self.episode_length_buf < self.max_episode_length - 1)).nonzero(as_tuple=False).flatten()
-        self._resample_commands(resampling_env_ids)
-        if self.cfg.commands.heading_command:
-            mask = (self.stop_heading == 0.0)
-            forward = quat_apply(self.base_quat[mask], self.forward_vec[mask])
-            heading = torch.atan2(forward[:, 1], forward[:, 0])
-            self.commands[mask, 2] = torch.clip(
-                0.5*wrap_to_pi(self.commands[mask, 3] - heading),
-                self.env_command_ranges["ang_vel_yaw"][:, 0],
-                self.env_command_ranges["ang_vel_yaw"][:, 1]
-            )
        if self.cfg.terrain.measure_heights:
            self.measured_heights = self._get_heights()

@@ -591,6 +581,17 @@ class LeggedRobot(BaseTask):

        self.commands_xy_accumulation[env_ids] += self.commands[env_ids, :2]

+        if self.cfg.commands.heading_command:
+            heading_env_ids = env_ids[self.stop_heading[env_ids] == 0.0]
+            if len(heading_env_ids) > 0:
+                forward = quat_apply(self.base_quat[heading_env_ids], self.forward_vec[heading_env_ids])
+                heading = torch.atan2(forward[:, 1], forward[:, 0])
+                self.commands[heading_env_ids, 2] = torch.clip(
+                    0.5 * wrap_to_pi(self.commands[heading_env_ids, 3] - heading),
+                    self.env_command_ranges["ang_vel_yaw"][heading_env_ids, 0],
+                    self.env_command_ranges["ang_vel_yaw"][heading_env_ids, 1]
+            )
+
    def _compute_torques(self, actions):
        """ Compute torques from actions.
            Actions can be interpreted as position or velocity targets given to a PD controller, or directly as scaled torques.
@@ -722,21 +723,7 @@ class LeggedRobot(BaseTask):
        self.gym.set_actor_root_state_tensor_indexed(self.sim,
                                                    gymtorch.unwrap_tensor(self.root_states),
                                                    gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32))
-
   
-    
-    def update_command_curriculum(self, env_ids):
-        """ Implements a curriculum of increasing commands
-
-        Args:
-            env_ids (List[int]): ids of environments being reset
-        """
-        # If the tracking reward is above 80% of the maximum, increase the range of commands
-        if torch.mean(self.episode_sums["tracking_lin_vel"][env_ids]) / self.max_episode_length > 0.8 * self.reward_scales["tracking_lin_vel"]:
-            self.command_ranges["lin_vel_x"][0] = np.clip(self.command_ranges["lin_vel_x"][0] - 0.5, -self.cfg.commands.max_curriculum, 0.)
-            self.command_ranges["lin_vel_x"][1] = np.clip(self.command_ranges["lin_vel_x"][1] + 0.5, 0., self.cfg.commands.max_curriculum)
-
-
    def _get_noise_scale_vec(self, cfg):
        """ Sets a vector used to scale the noise added to the observations.
            [NOTE]: Must be adapted when changing the observations structure
@@ -807,6 +794,7 @@ class LeggedRobot(BaseTask):
        self.d_gains = torch.zeros(self.num_actions, dtype=torch.float, device=self.device, requires_grad=False)
        self.actions = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False)
        self.last_actions = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False)
+        self.last_last_actions = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False)
        self.last_dof_vel = torch.zeros_like(self.dof_vel)
        self.last_root_vel = torch.zeros_like(self.root_states[:, 7:13])
        self.commands = torch.zeros(self.num_envs, self.cfg.commands.num_commands, dtype=torch.float, device=self.device, requires_grad=False) # x vel, y vel, yaw vel, heading
@@ -1372,10 +1360,7 @@ class LeggedRobot(BaseTask):

    def _reward_action_smoothness(self):
        # a_t - 2a_{t-1} + a_{t-2}
-        if not hasattr(self, 'last_last_actions'):
-            self.last_last_actions = torch.zeros_like(self.last_actions)
        rew = torch.sum((self.actions - 2 * self.last_actions + self.last_last_actions).pow(2), dim=1)
-        self.last_last_actions[:] = self.last_actions[:]
        return rew
    
    def _reward_dof_power(self):
--- a/legged_gym/envs/base/legged_robot_config.py
+++ b/legged_gym/envs/base/legged_robot_config.py
@@ -41,8 +41,6 @@ class LeggedRobotCfg(BaseConfig):
        move_down_by_accumulated_xy_command = False # move down the terrain curriculum based on accumulated xy command distance instead of absolute distance

    class commands:
-        curriculum = False
-        max_curriculum = 1.
        num_commands = 4 # default: lin_vel_x, lin_vel_y, ang_vel_yaw, heading (in heading mode ang_vel_yaw is recomputed from heading error)
        resampling_time = 10. # time before command are changed[s]
        heading_command = False # if true: compute ang vel command from heading error
--- a/legged_gym/envs/go2/go2_config.py
+++ b/legged_gym/envs/go2/go2_config.py
@@ -96,8 +96,6 @@ class GO2Cfg(LeggedRobotCfg):
        move_down_by_accumulated_xy_command = True # move down the terrain curriculum based on accumulated xy command distance instead of absolute distance
        
    class commands(LeggedRobotCfg.commands):
-        curriculum = False
-        max_curriculum = 1.
        num_commands = 4 # default: lin_vel_x, lin_vel_y, ang_vel_yaw (in heading mode ang_vel_yaw is recomputed from heading error)
        resampling_time = 5. # time before command are changed[s]
        heading_command = False # if true: compute ang vel command from heading error
@@ -151,7 +149,7 @@ class GO2Cfg(LeggedRobotCfg):
        foot_name = "foot"
        penalize_contacts_on = ["thigh", "calf"]
        terminate_after_contacts_on = ["base"]
-        self_collisions = 1 # 1 to disable, 0 to enable...bitwise filter
+        self_collisions = 0 # 1 to disable, 0 to enable...bitwise filter
  
    class rewards(LeggedRobotCfg.rewards):
        soft_dof_pos_limit = 0.9
--- a/legged_gym/envs/go2/go2_config_fast_flat_move.py
+++ b/legged_gym/envs/go2/go2_config_fast_flat_move.py
@@ -146,7 +146,7 @@ class GO2Cfg(LeggedRobotCfg):
            'heading': [-1.57, 1.57], # min max [rad]
        }, { # list for command range curriculums at specific training iterations
            'iter': 40000, # training iteration at which the command ranges are updated
-            'lin_vel_x': [-2.0, 4.2], # min max [m/s]
+            'lin_vel_x': [-2.0, 4.5], # min max [m/s]
            'lin_vel_y': [-0.5, 0.5], # min max [m/s]
            'ang_vel_yaw': [-1.0, 1.0], # min max [rad/s]
            'heading': [-1.57, 1.57], # min max [rad]
@@ -181,7 +181,7 @@ class GO2Cfg(LeggedRobotCfg):
        foot_name = "foot"
        penalize_contacts_on = ["thigh", "calf"]
        terminate_after_contacts_on = ["base"]
-        self_collisions = 1 # 1 to disable, 0 to enable...bitwise filter
+        self_collisions = 0 # 1 to disable, 0 to enable...bitwise filter
  
    class rewards(LeggedRobotCfg.rewards):
        soft_dof_pos_limit = 0.9
--- a/legged_gym/envs/go2/go2_config_vanilla.py
+++ b/legged_gym/envs/go2/go2_config_vanilla.py
@@ -165,7 +165,7 @@ class GO2Cfg(LeggedRobotCfg):
        foot_name = "foot"
        penalize_contacts_on = ["thigh", "calf"]
        terminate_after_contacts_on = ["base"]
-        self_collisions = 1 # 1 to disable, 0 to enable...bitwise filter
+        self_collisions = 0 # 1 to disable, 0 to enable...bitwise filter
  
    class rewards(LeggedRobotCfg.rewards):
        soft_dof_pos_limit = 0.9
--- a/legged_gym/envs/go2/go2_config_vanilla_with_dynamic_cmd.py
+++ b/legged_gym/envs/go2/go2_config_vanilla_with_dynamic_cmd.py
@@ -165,7 +165,7 @@ class GO2Cfg(LeggedRobotCfg):
        foot_name = "foot"
        penalize_contacts_on = ["thigh", "calf"]
        terminate_after_contacts_on = ["base"]
-        self_collisions = 1 # 1 to disable, 0 to enable...bitwise filter
+        self_collisions = 0 # 1 to disable, 0 to enable...bitwise filter
  
    class rewards(LeggedRobotCfg.rewards):
        soft_dof_pos_limit = 0.9
--- a/rsl_rl/rsl_rl/runners/on_policy_runner.py
+++ b/rsl_rl/rsl_rl/runners/on_policy_runner.py
@@ -253,38 +253,56 @@ class OnPolicyRunner:
        if self.robogauge_client is None:
            return

-        if it % 500 == 0 or last_model:
-            # export jit model
-            jit_dir = os.path.join(self.log_dir, 'jit_models')
-            jit_path = os.path.join(jit_dir, f'policy_jit_{it}.pt')
-            export_policy_as_jit(self.alg.actor_critic, jit_dir, filename=f'policy_jit_{it}.pt')
-            # upload to robogauge
-            task_name = 'go2'
-            self.robogauge_client.submit_task(
-                model_path=jit_path,
-                step=it,
-                task_name=task_name,
-                experiment_name=self.cfg["experiment_name"]
-            )
+        try:
+            if it % 500 == 0 or last_model:
+                # export jit model
+                jit_dir = os.path.join(self.log_dir, 'jit_models')
+                jit_path = os.path.join(jit_dir, f'policy_jit_{it}.pt')
+                export_policy_as_jit(self.alg.actor_critic, jit_dir, filename=f'policy_jit_{it}.pt')
+                # upload to robogauge
+                task_name = 'go2'
+                self.robogauge_client.submit_task(
+                    model_path=jit_path,
+                    step=it,
+                    task_name=task_name,
+                    experiment_name=self.cfg["experiment_name"]
+                )
+        except Exception as e:
+            print(f"[WARN] RoboGauge submit failed at step {it}: {e}")
+            return
        check_times = 1
        if last_model:
            check_times = int(1e9)  # keep checking until the last model is evaluated
        while check_times > 0:
            check_times -= 1
-            self.robogauge_client.monitor_tasks()
+            try:
+                self.robogauge_client.monitor_tasks()
+            except Exception as e:
+                print(f"[WARN] RoboGauge monitor failed at step {it}: {e}")
+                break
            results_dir = os.path.join(self.log_dir, 'robogauge_results')
            os.makedirs(results_dir, exist_ok=True)
            result_received = False
            for task_id, resp in self.robogauge_client.response_data.items():
-                scores = resp['results']['scores']
-                step = resp['step']
+                if not isinstance(resp, dict):
+                    print(f"[WARN] RoboGauge returned an invalid response for task {task_id}: {resp}")
+                    continue
+                results = resp.get('results')
+                step = resp.get('step', it)
+                if results is None:
+                    print(f"[WARN] RoboGauge returned empty results for task {task_id} at step {step}.")
+                    continue
+                scores = results.get('scores')
+                if scores is None:
+                    print(f"[WARN] RoboGauge results for task {task_id} at step {step} do not contain 'scores'.")
+                    continue
                if step == it:
                    result_received = True
                for key, val in scores.items():
                    self.writer.add_scalar(f'RoboGauge/{key}', val, step)
                results_path = os.path.join(results_dir, f'results_{step}.yaml')
                with open(results_path, 'w', encoding='utf-8') as f:
-                    yaml.dump(resp['results'], f, allow_unicode=True, sort_keys=False)
+                    yaml.dump(results, f, allow_unicode=True, sort_keys=False)
            
            if last_model and result_received:
                print(f"RoboGauge result for step {it} received. Exiting wait loop.")
--- a/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py
+++ b/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py
@@ -298,38 +298,56 @@ class OnPolicyRunnerCTS:
        if self.robogauge_client is None:
            return

-        if it % 500 == 0 or last_model:
-            # export jit model
-            jit_dir = os.path.join(self.log_dir, 'jit_models')
-            jit_path = os.path.join(jit_dir, f'policy_jit_{it}.pt')
-            export_policy_as_jit(self.alg.model, jit_dir, filename=f'policy_jit_{it}.pt')
-            # upload to robogauge
-            task_name = 'go2_moe'  # Both cts, moe-cts actor return a tuple `action, (latent, ...)`
-            self.robogauge_client.submit_task(
-                model_path=jit_path,
-                step=it,
-                task_name=task_name,
-                experiment_name=self.cfg["experiment_name"]
-            )
+        try:
+            if it % 500 == 0 or last_model:
+                # export jit model
+                jit_dir = os.path.join(self.log_dir, 'jit_models')
+                jit_path = os.path.join(jit_dir, f'policy_jit_{it}.pt')
+                export_policy_as_jit(self.alg.model, jit_dir, filename=f'policy_jit_{it}.pt')
+                # upload to robogauge
+                task_name = 'go2_moe'  # Both cts, moe-cts actor return a tuple `action, (latent, ...)`
+                self.robogauge_client.submit_task(
+                    model_path=jit_path,
+                    step=it,
+                    task_name=task_name,
+                    experiment_name=self.cfg["experiment_name"]
+                )
+        except Exception as e:
+            print(f"[WARN] RoboGauge submit failed at step {it}: {e}")
+            return
        check_times = 1
        if last_model:
            check_times = int(1e9)  # keep checking until manually stopped
        while check_times > 0:
            check_times -= 1
-            self.robogauge_client.monitor_tasks()
+            try:
+                self.robogauge_client.monitor_tasks()
+            except Exception as e:
+                print(f"[WARN] RoboGauge monitor failed at step {it}: {e}")
+                break
            results_dir = os.path.join(self.log_dir, 'robogauge_results')
            os.makedirs(results_dir, exist_ok=True)
            result_received = False
            for task_id, resp in self.robogauge_client.response_data.items():
-                scores = resp['results']['scores']
-                step = resp['step']
+                if not isinstance(resp, dict):
+                    print(f"[WARN] RoboGauge returned an invalid response for task {task_id}: {resp}")
+                    continue
+                results = resp.get('results')
+                step = resp.get('step', it)
+                if results is None:
+                    print(f"[WARN] RoboGauge returned empty results for task {task_id} at step {step}.")
+                    continue
+                scores = results.get('scores')
+                if scores is None:
+                    print(f"[WARN] RoboGauge results for task {task_id} at step {step} do not contain 'scores'.")
+                    continue
                if step == it:
                    result_received = True
                for key, val in scores.items():
                    self.writer.add_scalar(f'RoboGauge/{key}', val, step)
                results_path = os.path.join(results_dir, f'results_{step}.yaml')
                with open(results_path, 'w', encoding='utf-8') as f:
-                    yaml.dump(resp['results'], f, allow_unicode=True, sort_keys=False)
+                    yaml.dump(results, f, allow_unicode=True, sort_keys=False)
            
            if last_model and result_received:
                print(f"RoboGauge result for step {it} received. Exiting wait loop.")
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from setuptools import find_packages
 from distutils.core import setup

 setup(name='go2_rl_gym',
-      version='1.0.2',
+      version='1.0.3',
      author='Wu Tianyang',
      license="MIT",
      packages=find_packages(),
--- a/tools/logs_merge.py
+++ b/tools/logs_merge.py
@@ -36,7 +36,39 @@ def fast_read(event_file_path, tag_names):
                if value.tag in tag_names:
                    tag_data[event.step][value.tag] = value.simple_value
                    
-    return pd.DataFrame(tag_data).T
+    df = pd.DataFrame(tag_data).T
+    df.index.name = 'step'
+    return df
+
+def normalize_tb_df(tb_df):
+    tb_df = tb_df.copy()
+
+    if 'step' not in tb_df.columns:
+        first_col = tb_df.columns[0] if len(tb_df.columns) > 0 else None
+        if first_col is not None and str(first_col).startswith('Unnamed:'):
+            tb_df = tb_df.rename(columns={first_col: 'step'})
+        elif tb_df.index.name == 'step':
+            tb_df = tb_df.reset_index()
+        else:
+            tb_df = tb_df.reset_index().rename(columns={'index': 'step'})
+
+    tb_df['step'] = pd.to_numeric(tb_df['step'], errors='coerce')
+    tb_df = tb_df.dropna(subset=['step'])
+    tb_df['step'] = tb_df['step'].astype(int)
+    return tb_df
+
+def get_tb_value(tb_df, step, candidate_tags):
+    row = tb_df[tb_df['step'] == step]
+    if row.empty:
+        raise KeyError(f"No tensorboard entry found for step={step}.")
+
+    for tag in candidate_tags:
+        if tag not in row.columns:
+            continue
+        values = row[tag].dropna().values
+        if len(values) > 0:
+            return float(values[0])
+    raise KeyError(f"No tensorboard value found for step={step} in tags: {candidate_tags}")

 class Collector:
    def __init__(self, log_dirs):
@@ -58,14 +90,14 @@ class Collector:
        self.output_tb = self.output_dir / "tb.csv"
        if self.output_tb.exists():
            print(f"Loading existing tensorboard data from {self.output_tb}")
-            self.tb_df = pd.read_csv(self.output_tb)
+            self.tb_df = normalize_tb_df(pd.read_csv(self.output_tb))
        else:
            start_time = time.time()
            print(f"Start reading tensorboard events at {time.ctime(start_time)}")
-            self.tb_df = fast_read(str(self.log_dirs.glob("events.out.tfevents.*").__next__()), [
+            self.tb_df = normalize_tb_df(fast_read(str(self.log_dirs.glob("events.out.tfevents.*").__next__()), [
                'Terrain/terrain_level_all', 'Episode/terrain_level_all',
                'RoboGauge/benchmark'
-            ])
+            ]))
            print(f"Finished reading tensorboard events in {time.time() - start_time:.2f} seconds.")
            self.tb_df.to_csv(self.output_tb, index=False)
            print(f"Saved tensorboard data to {self.output_tb}")
@@ -109,7 +141,11 @@ class Collector:
                self.datas[f'{terrain_name}_mean@25'].append(float(data['robust_score'][terrain_name]['mean@25']))
                self.datas[f'{terrain_name}_mean@50'].append(float(data['robust_score'][terrain_name]['mean@50']))
            
-            self.datas['terrain_level'].append(float(self.tb_df[self.tb_df['step'] == it]['value'].values[0]))
+            self.datas['terrain_level'].append(get_tb_value(
+                self.tb_df,
+                it,
+                ['Terrain/terrain_level_all', 'Episode/terrain_level_all']
+            ))
        df = pd.DataFrame(self.datas)
        df.to_csv(self.output_csv, index=False)
        print(f"Saved merged results to {self.output_csv}")
@@ -117,6 +153,8 @@ class Collector:
 if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--log-dirs")
+    parser.add_argument("--read-robogauge", default=True, type=lambda x: (str(x).lower() in ['true', '1']), help="Whether to read robogauge_results")
    args = parser.parse_args()
    collector = Collector(args.log_dirs)
-    # collector.collect()
+    if args.read_robogauge:
+        collector.collect()