v0.1.2; add ppo robogauge

UPDATE.md

@@ -1,3 +1,6 @@
+# 20251231
+## v0.1.2
+1. 加入PPO的Robogauge评估
 # 20251230
 ## v0.1.1
 1. 给cts算法加入robogauge异步评估

legged_gym/utils/exporter.py

@@ -183,6 +183,9 @@ class _OnnxPolicyExporter(torch.nn.Module):
             self.history_length = policy.history.shape[1]
             self.forward = self.forward_moe_cts
             self.input_dim = self.history_length * policy.history.shape[2]
+        
+        else:  # PPO
+            self.forward = self.forward_ppo
             
         if hasattr(policy, "actor"):
             self.actor = copy.deepcopy(policy.actor)
@@ -198,8 +201,7 @@ class _OnnxPolicyExporter(torch.nn.Module):
         else:
             raise ValueError("Policy does not have an actor/student module.")
 
-    def forward_cts(self, x):  # x is stack observations by terms
+    def flatten_obs(self, x):  # flatten stack obs by terms to stack by step frames
-        x = self.normalizer(x)
         term_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
         obs_dim = sum(term_dims)
         if x.shape[1] % obs_dim != 0:
@@ -223,6 +225,17 @@ class _OnnxPolicyExporter(torch.nn.Module):
             history_by_frame.append(torch.cat(terms_for_this_frame, dim=1))
         # [B, (Frame0_AllTerms), (Frame1_AllTerms), ...]
         history = torch.cat(history_by_frame, dim=1)
+        return history, obs_dim
+    
+    def forward_ppo(self, x):  # x is stack observations by terms
+        x = self.normalizer(x)
+        history, obs_dim = self.flatten_obs(x)
+        last_obs = history[:, -obs_dim:]
+        return self.actor(last_obs)
+
+    def forward_cts(self, x):  # x is stack observations by terms
+        x = self.normalizer(x)
+        history, obs_dim = self.flatten_obs(x)
 
         last_obs = history[:, -obs_dim:]
         latent = self.student_encoder(history)
@@ -232,29 +245,7 @@ class _OnnxPolicyExporter(torch.nn.Module):
 
     def forward_moe_cts(self, x):
         x = self.normalizer(x)
-        term_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
+        history, obs_dim = self.flatten_obs(x)
-        obs_dim = sum(term_dims)
-        if x.shape[1] % obs_dim != 0:
-            raise ValueError(f"x.shape[1] ({x.shape[1]}) 不是 obs_dim ({obs_dim}) 的整数倍")
-            
-        frames = x.shape[1] // obs_dim
-        split_sizes = [dim * frames for dim in term_dims]
-        # [B, dim0*frames], [B, dim1*frames], ...
-        term_chunks = torch.split(x, split_sizes, dim=1)
-
-        # [ [B, frames, dim0], [B, frames, dim1], ... ]
-        frame_terms_reshaped = [
-            chunk.view(-1, frames, dim) 
-            for chunk, dim in zip(term_chunks, term_dims)
-        ]
-
-        history_by_frame = []
-        for i in range(frames):
-            # [ [B, dim0], [B, dim1], ... ]
-            terms_for_this_frame = [ftr[:, i, :] for ftr in frame_terms_reshaped]
-            history_by_frame.append(torch.cat(terms_for_this_frame, dim=1))
-        # [B, (Frame0_AllTerms), (Frame1_AllTerms), ...]
-        history = torch.cat(history_by_frame, dim=1)
 
         last_obs = history[:, -obs_dim:]
         history_3d = history.view(-1, self.history_length, obs_dim)
@@ -267,25 +258,11 @@ class _OnnxPolicyExporter(torch.nn.Module):
     
     def forward_mcp_cts(self, x):
         x = self.normalizer(x)
-        term_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
+        history, obs_dim = self.flatten_obs(x)
-        obs_dim = sum(term_dims)
-        frames = x.shape[1] // obs_dim
-        
-        split_sizes = [dim * frames for dim in term_dims]
-        term_chunks = torch.split(x, split_sizes, dim=1)
-        frame_terms_reshaped = [chunk.view(-1, frames, dim) for chunk, dim in zip(term_chunks, term_dims)]
-        history_by_frame = []
-        for i in range(frames):
-            terms_for_this_frame = [ftr[:, i, :] for ftr in frame_terms_reshaped]
-            history_by_frame.append(torch.cat(terms_for_this_frame, dim=1))
-        history = torch.cat(history_by_frame, dim=1)
 
         last_obs = history[:, -obs_dim:]
-        
         obs_no_goal = last_obs[:, self.obs_no_goal_mask]
-        
         latent = self.student_encoder(history)
-        
         x_in = torch.cat([latent, last_obs], dim=1)
         x_no_goal_in = torch.cat([latent, obs_no_goal], dim=1)
         

rsl_rl/rsl_rl/runners/on_policy_runner.py

@@ -32,6 +32,7 @@ import time
 import os
 from collections import deque
 import statistics
+import yaml
 
 from torch.utils.tensorboard import SummaryWriter
 import torch
@@ -39,6 +40,7 @@ import torch
 from rsl_rl.algorithms import PPO
 from rsl_rl.modules import ActorCritic, ActorCriticRecurrent
 from rsl_rl.env import VecEnv
+from legged_gym.utils.exporter import export_policy_as_jit
 
 
 class OnPolicyRunner:
@@ -79,6 +81,13 @@ class OnPolicyRunner:
         self.current_learning_iteration = 0
 
         _, _ = self.env.reset()
+
+        # robogauge client
+        try:
+            from robogauge.scripts.client import RoboGaugeClient
+            self.robogauge_client = RoboGaugeClient()
+        except:
+            self.robogauge_client = None
     
     def learn(self, num_learning_iterations, init_at_random_ep_len=False):
         # initialize writer
@@ -135,11 +144,11 @@ class OnPolicyRunner:
             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.current_learning_iteration += num_learning_iterations
-        self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(self.current_learning_iteration)))
+        self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(self.current_learning_iteration)), it)
 
     def log(self, locs, width=80, pad=35):
         self.tot_timesteps += self.num_steps_per_env * self.env.num_envs
@@ -210,13 +219,43 @@ class OnPolicyRunner:
                                locs['num_learning_iterations'] - 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.actor_critic.state_dict(),
             'optimizer_state_dict': self.alg.optimizer.state_dict(),
             'iter': self.current_learning_iteration,
             'infos': infos,
             }, path)
+        self.update_robogauge(it)
+
+    def update_robogauge(self, it):
+        if self.robogauge_client is None:
+            return
+
+        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.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"]
+            )
+        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)

rsl_rl/rsl_rl/runners/on_policy_runner_cts.py

@@ -45,7 +45,6 @@ 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):
@@ -112,7 +111,11 @@ class OnPolicyRunnerCTS:
             yaml.safe_dump(all_cfg, open(os.path.join(self.log_dir, 'config.yaml'), 'w'))
         
         # robogauge client
-        self.robogauge_client = RoboGaugeClient()
+        try:
+            from robogauge.scripts.client import RoboGaugeClient
+            self.robogauge_client = RoboGaugeClient()
+        except:
+            self.robogauge_client = None
     
     def learn(self, num_learning_iterations, init_at_random_ep_len=False):
         # initialize writer
@@ -188,7 +191,7 @@ class OnPolicyRunnerCTS:
                 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)))
+        self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(self.current_learning_iteration)), it)
 
     def log(self, locs, width=80, pad=35):
         self.tot_timesteps += self.num_steps_per_env * self.env.num_envs
@@ -280,6 +283,9 @@ class OnPolicyRunnerCTS:
         self.update_robogauge(it)
     
     def update_robogauge(self, it):
+        if self.robogauge_client is None:
+            return
+
         if it % 500 == 0:
             # export jit model
             jit_dir = os.path.join(self.log_dir, 'jit_models')