v0.1.1; Add robogauge to CTS Runner

UPDATE.md

@@ -1,3 +1,6 @@
+# 20251230
+## v0.1.1
+1. 给cts算法加入robogauge异步评估
 # 20251221
 1. 修改最大地形速度限制, y在所有地形上最大为1.0, z只有平地最大为2.0, x最大为2.0
 2. 上调难度9地形难度 (都是moe-cts 100k能通过的难度):

doc/setup_zh.md

@@ -88,6 +88,12 @@ python 1080_balls_of_solitude.py
 
 `rsl_rl` 是一个强化学习算法库。
 
+我们仓库中是带有新算法的 `rsl_rl`，克隆 Git 仓库：
+
+```bash
+git clone https://github.com/wty-yy/go2_rl_gym.git
+```
+
 #### 2.3.1 安装
 
 ```bash
@@ -97,16 +103,6 @@ pip install -e .
 
 ### 2.4 安装 go2_rl_gym
 
-#### 2.4.1 下载
-
-通过 Git 克隆仓库：
-
-```bash
-git clone https://github.com/wty-yy/go2_rl_gym.git
-```
-
-#### 2.4.2 安装
-
 进入目录并安装：
 
 ```bash

legged_gym/envs/go2/go2_config.py

@@ -244,7 +244,7 @@ class GO2CfgCTS(LeggedRobotCfgCTS):
         num_steps_per_env = 24
         run_name = ''
         experiment_name = 'go2_cts'
-        max_iterations = 100000
+        max_iterations = 150000
         save_interval = 500
     
     class policy(LeggedRobotCfgCTS.policy):
@@ -262,5 +262,5 @@ class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
     class runner(LeggedRobotCfgMoECTS.runner):
         run_name = ''
         experiment_name = 'go2_moe_cts'
-        max_iterations = 100000
+        max_iterations = 150000
         save_interval = 500

rsl_rl/rsl_rl/runners/on_policy_runner_cts.py

@@ -45,6 +45,8 @@ import numpy as np
 from pathlib import Path
 from legged_gym.utils.helpers import class_to_dict
 from typing import Union
+from robogauge.scripts.client import RoboGaugeClient
+from legged_gym.utils.exporter import export_policy_as_jit
 
 def numpy_representer(dumper, data):
     return dumper.represent_float(float(data))
@@ -109,6 +111,9 @@ class OnPolicyRunnerCTS:
             all_cfg = {"train_cfg": train_cfg, "env_cfg": class_to_dict(self.env.cfg)}
             yaml.safe_dump(all_cfg, open(os.path.join(self.log_dir, 'config.yaml'), 'w'))
         
+        # robogauge client
+        self.robogauge_client = RoboGaugeClient()
+    
     def learn(self, num_learning_iterations, init_at_random_ep_len=False):
         # initialize writer
         if self.log_dir is not None and self.writer is None:
@@ -180,7 +185,7 @@ class OnPolicyRunnerCTS:
             if self.log_dir is not None:
                 self.log(locals())
             if it % self.save_interval == 0:
-                self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(it)))
+                self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(it)), it)
             ep_infos.clear()
         
         self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(self.current_learning_iteration)))
@@ -261,7 +266,7 @@ class OnPolicyRunnerCTS:
                                locs['tot_iter'] - locs['it']):.1f}s\n""")
         print(log_string)
 
-    def save(self, path, infos=None):
+    def save(self, path, it, infos=None):
         torch.save({
             'model_state_dict': self.alg.model.state_dict(),
             'optimizer1_state_dict': self.alg.optimizer1.state_dict(),
@@ -269,6 +274,32 @@ class OnPolicyRunnerCTS:
             'iter': self.current_learning_iteration,
             'infos': infos,
             }, path)
+        self.update_robogauge(path, it)
+    
+    def update_robogauge(self, model_path, it):
+        if it % 500 == 0:
+            # 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
+            self.robogauge_client.submit_task(
+                model_path=jit_path,
+                step=it,
+                task_name='go2_moe' if 'moe' in self.cfg["algorithm_class_name"].lower() else 'go2',
+                experiment_name=self.cfg["experiment_name"]
+            )
+        self.robogauge_client.monitor_tasks()
+        results_dir = os.path.join(self.log_dir, 'robogauge_results')
+        os.makedirs(results_dir, exist_ok=True)
+        for task_id, resp in self.robogauge_client.response_data.items():
+            scores = resp['results']['scores']
+            step = resp['step']
+            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)
 
     def load(self, path, load_optimizer=True):
         loaded_dict = torch.load(path)