cyy/go2_rl_gym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

v0.1.2; add ppo robogauge

Browse Source
This commit is contained in:
wty-yy
2025-12-31 23:16:52 +08:00
parent 9ed3f0e144
commit 102c0f8df7
4 changed files with 71 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
UPDATE.md
View File

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

57
legged_gym/utils/exporter.py
View File

@@ -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)

45
rsl_rl/rsl_rl/runners/on_policy_runner.py
View File

@@ -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)

12
rsl_rl/rsl_rl/runners/on_policy_runner_cts.py
View File

@@ -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')