From 102c0f8df7b1d33ee67f2adacebc264ecd65958e Mon Sep 17 00:00:00 2001 From: wty-yy <993660140@qq.com> Date: Wed, 31 Dec 2025 23:16:52 +0800 Subject: [PATCH] v0.1.2; add ppo robogauge --- UPDATE.md | 3 + legged_gym/utils/exporter.py | 57 ++++++------------- rsl_rl/rsl_rl/runners/on_policy_runner.py | 45 ++++++++++++++- rsl_rl/rsl_rl/runners/on_policy_runner_cts.py | 12 +++- 4 files changed, 71 insertions(+), 46 deletions(-) diff --git a/UPDATE.md b/UPDATE.md index 2b0acb9..3c4c25d 100644 --- a/UPDATE.md +++ b/UPDATE.md @@ -1,3 +1,6 @@ +# 20251231 +## v0.1.2 +1. 加入PPO的Robogauge评估 # 20251230 ## v0.1.1 1. 给cts算法加入robogauge异步评估 diff --git a/legged_gym/utils/exporter.py b/legged_gym/utils/exporter.py index 04339ff..b590415 100644 --- a/legged_gym/utils/exporter.py +++ b/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 - x = self.normalizer(x) + def flatten_obs(self, x): # flatten stack obs by terms to stack by step frames 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] - 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) + history, obs_dim = self.flatten_obs(x) 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] - 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) + history, obs_dim = self.flatten_obs(x) 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) diff --git a/rsl_rl/rsl_rl/runners/on_policy_runner.py b/rsl_rl/rsl_rl/runners/on_policy_runner.py index e11ae4d..a808a2f 100644 --- a/rsl_rl/rsl_rl/runners/on_policy_runner.py +++ b/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) diff --git a/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py b/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py index bf9c473..38a8032 100644 --- a/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py +++ b/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')