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

Add mcp-cts

Browse Source
This commit is contained in:
wty-yy
2025-12-31 01:16:33 +08:00
parent 010c5b1700
commit 9ed3f0e144
11 changed files with 609 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
legged_gym/envs/__init__.py
View File

@@ -1,11 +1,12 @@
from legged_gym import LEGGED_GYM_ROOT_DIR, LEGGED_GYM_ENVS_DIR
from legged_gym.envs.go2.go2_env import Go2Robot
from legged_gym.envs.go2.go2_config import GO2Cfg, GO2CfgPPO, GO2CfgCTS, GO2CfgMoECTS
from legged_gym.envs.go2.go2_config import GO2Cfg, GO2CfgPPO, GO2CfgCTS, GO2CfgMoECTS, GO2CfgMCPCTS
from .base.legged_robot import LeggedRobot
from legged_gym.utils.task_registry import task_registry
task_registry.register( "go2", Go2Robot, GO2Cfg(), GO2CfgPPO())
task_registry.register( "go2_cts", Go2Robot, GO2Cfg(), GO2CfgCTS())
task_registry.register( "go2_moe_cts", Go2Robot, GO2Cfg(), GO2CfgMoECTS())
task_registry.register("go2", Go2Robot, GO2Cfg(), GO2CfgPPO())
task_registry.register("go2_cts", Go2Robot, GO2Cfg(), GO2CfgCTS())
task_registry.register("go2_moe_cts", Go2Robot, GO2Cfg(), GO2CfgMoECTS())
task_registry.register("go2_mcp_cts", Go2Robot, GO2Cfg(), GO2CfgMCPCTS())

9
legged_gym/envs/base/legged_robot_config.py
View File

@@ -359,3 +359,12 @@ class LeggedRobotCfgMoECTS(LeggedRobotCfgCTS):
class runner(LeggedRobotCfgCTS.runner):
policy_class_name = 'ActorCriticMoECTS'
algorithm_class_name = 'MoECTS'
class LeggedRobotCfgMCPCTS(LeggedRobotCfgCTS):
class policy(LeggedRobotCfgCTS.policy):
obs_no_goal_mask = None # mask for observation without goal inputs
student_expert_num = 8 # number of experts in the student model
class runner(LeggedRobotCfgCTS.runner):
policy_class_name = 'ActorCriticMCPCTS'
algorithm_class_name = 'MCPCTS'

13
legged_gym/envs/go2/go2_config.py
View File

@@ -1,5 +1,5 @@
import math
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMCPCTS
class GO2Cfg(LeggedRobotCfg):
class init_state(LeggedRobotCfg.init_state):
@@ -264,3 +264,14 @@ class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
experiment_name = 'go2_moe_cts'
max_iterations = 150000
save_interval = 500
class GO2CfgMCPCTS(LeggedRobotCfgMCPCTS):
class policy(LeggedRobotCfgMCPCTS.policy):
obs_no_goal_mask = [True] * 6 + [False] * 3 + [True] * 36 # mask for obs without command info
student_expert_num = 8 # number of experts in the student model
class runner(LeggedRobotCfgMCPCTS.runner):
run_name = ''
experiment_name = 'go2_mcp_cts'
max_iterations = 150000
save_interval = 500

50
legged_gym/utils/exporter.py
View File

@@ -81,7 +81,11 @@ class _TorchPolicyExporter(torch.nn.Module):
self.history_length = policy.history.shape[1]
self.history = torch.zeros([1, policy.history.shape[1], policy.history.shape[2]], device='cpu')
self.forward = self.forward_moe_cts
if hasattr(policy, "actor"):
if hasattr(policy, "actor_mcp"):
self.actor = copy.deepcopy(policy.actor_mcp)
self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
self.forward = self.forward_mcp_cts
elif hasattr(policy, "actor"):
self.actor = copy.deepcopy(policy.actor)
if self.is_recurrent:
self.rnn = copy.deepcopy(policy.memory_a.rnn)
@@ -129,6 +133,16 @@ class _TorchPolicyExporter(torch.nn.Module):
latent, weights = self.student_moe_encoder(self.history.flatten(1), history_no_goal)
x = torch.cat([latent, x], dim=1)
return self.actor(x), (weights, latent)
def forward_mcp_cts(self, x): # x is single observations
x = self.normalizer(x)
self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
x_no_goal = x[:, self.obs_no_goal_mask]
latent = self.student_encoder(self.history.flatten(1))
x = torch.cat([latent, x], dim=1)
x_no_goal = torch.cat([latent, x_no_goal], dim=1)
mean_action, _, weights = self.actor(x, x_no_goal)
return mean_action, weights
@torch.jit.export
def reset(self):
@@ -176,6 +190,11 @@ class _OnnxPolicyExporter(torch.nn.Module):
self.rnn = copy.deepcopy(policy.memory_a.rnn)
if self.input_dim is None:
self.input_dim = self.actor[0].in_features
elif hasattr(policy, "actor_mcp"):
self.actor = copy.deepcopy(policy.actor_mcp)
self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
self.history_length = policy.history.shape[1]
self.forward = self.forward_mcp_cts
else:
raise ValueError("Policy does not have an actor/student module.")
@@ -245,6 +264,33 @@ class _OnnxPolicyExporter(torch.nn.Module):
x = torch.cat([latent, last_obs], dim=1)
return self.actor(x), weights, latent
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)
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)
mean_action, _, weights = self.actor(x_in, x_no_goal_in)
return mean_action, weights
def export(self, path, filename):
self.to("cpu")
@@ -254,6 +300,8 @@ class _OnnxPolicyExporter(torch.nn.Module):
if self.forward == self.forward_moe_cts:
output_names.append("weights")
output_names.append("latent")
if self.forward == self.forward_mcp_cts:
output_names.append("weights")
torch.onnx.export(
self,