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

v0.1.5 prev1; Add DualMoE

Browse Source
This commit is contained in:
wty-yy
2026-01-07 00:21:59 +08:00
parent 05e1e81d64
commit e4aa714eab
10 changed files with 487 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
legged_gym/envs/__init__.py
View File

@@ -1,7 +1,7 @@
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, GO2CfgMCPCTS, GO2CfgACMoECTS
from legged_gym.envs.go2.go2_config import GO2Cfg, GO2CfgPPO, GO2CfgCTS, GO2CfgMoECTS, GO2CfgMCPCTS, GO2CfgACMoECTS, GO2CfgDualMoECTS
from .base.legged_robot import LeggedRobot
from legged_gym.utils.task_registry import task_registry
@@ -11,3 +11,4 @@ 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())
task_registry.register("go2_ac_moe_cts", Go2Robot, GO2Cfg(), GO2CfgACMoECTS())
task_registry.register("go2_dual_moe_cts", Go2Robot, GO2Cfg(), GO2CfgDualMoECTS())

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

@@ -375,4 +375,13 @@ class LeggedRobotCfgACMoECTS(LeggedRobotCfgCTS):
class runner(LeggedRobotCfgCTS.runner):
policy_class_name = 'ActorCriticACMoECTS'
algorithm_class_name = 'ACMoECTS'
algorithm_class_name = 'ACMoECTS'
class LeggedRobotCfgDualMoECTS(LeggedRobotCfgCTS):
class policy(LeggedRobotCfgCTS.policy):
expert_num = 8 # number of experts in the student model
student_encoder_hidden_dims = [512, 256, 128]
class runner(LeggedRobotCfgCTS.runner):
policy_class_name = 'ActorCriticDualMoECTS'
algorithm_class_name = 'DualMoECTS'

12
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, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS, LeggedRobotCfgDualMoECTS
class GO2Cfg(LeggedRobotCfg):
class init_state(LeggedRobotCfg.init_state):
@@ -285,3 +285,13 @@ class GO2CfgACMoECTS(LeggedRobotCfgACMoECTS):
experiment_name = 'go2_ac_moe_cts'
max_iterations = 150000
save_interval = 500
class GO2CfgDualMoECTS(LeggedRobotCfgDualMoECTS):
class policy(LeggedRobotCfgDualMoECTS.policy):
expert_num = 8 # number of experts in the student model
class runner(LeggedRobotCfgDualMoECTS.runner):
run_name = ''
experiment_name = 'go2_dual_moe_cts'
max_iterations = 150000
save_interval = 500

13
legged_gym/utils/exporter.py
View File

@@ -77,7 +77,8 @@ class _TorchPolicyExporter(torch.nn.Module):
self.forward = self.forward_cts
if hasattr(policy, "student_moe_encoder"):
self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder).cpu()
self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
if hasattr(policy, "obs_no_goal_mask"):
self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
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
@@ -98,6 +99,8 @@ class _TorchPolicyExporter(torch.nn.Module):
self.rnn = copy.deepcopy(policy.memory_s.rnn)
else:
raise ValueError("Policy does not have an actor/student module.")
if hasattr(policy, "student_moe_encoder") and hasattr(policy, "actor_moe"):
self.forward = self.forward_dual_moe_cts
# set up recurrent network
if self.is_recurrent:
self.rnn.cpu()
@@ -155,6 +158,14 @@ class _TorchPolicyExporter(torch.nn.Module):
mean, weights = self.actor(x)
return mean, (weights, latent)
def forward_dual_moe_cts(self, x): # x is single observations
x = self.normalizer(x)
self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
latent, student_weights = self.student_moe_encoder(self.history.flatten(1))
x = torch.cat([latent, x], dim=1)
mean, actor_weights = self.actor(x)
return mean, (student_weights, actor_weights, latent)
@torch.jit.export
def reset(self):
if hasattr(self, 'history'):