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

v1.0.2-rc1; mv rem_cts to moe_cts, moe_cts to moe_no_goal_cts, fix rem params same as moe, add --robogauge to start

Browse Source
This commit is contained in:
wty-yy
2026-01-26 17:48:34 +08:00
parent 2aed91e7ce
commit 1798e67c29
23 changed files with 438 additions and 409 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
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, GO2CfgDualMoECTS, GO2CfgREMCTS
from legged_gym.envs.go2.go2_config import GO2Cfg, GO2CfgPPO, GO2CfgCTS, GO2CfgMoECTS, GO2CfgMoENGCTS, GO2CfgMCPCTS, GO2CfgACMoECTS, GO2CfgDualMoECTS
from .base.legged_robot import LeggedRobot
from legged_gym.utils.task_registry import task_registry
@@ -9,7 +9,7 @@ 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_moe_ng_cts", Go2Robot, GO2Cfg(), GO2CfgMoENGCTS())
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())
task_registry.register("go2_rem_cts", Go2Robot, GO2Cfg(), GO2CfgREMCTS())

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

@@ -302,6 +302,10 @@ class LeggedRobotCfgPPO(BaseConfig):
checkpoint = -1 # -1 = last saved model
resume_path = None # updated from load_run and chkpt
class robogauge:
enabled = False
port = 9973
class LeggedRobotCfgCTS(BaseConfig):
seed = 0
runner_class_name = "OnPolicyRunnerCTS"
@@ -350,7 +354,11 @@ class LeggedRobotCfgCTS(BaseConfig):
checkpoint = -1 # -1 = last saved model
resume_path = None # updated from load_run and chkpt
class LeggedRobotCfgMoECTS(LeggedRobotCfgCTS):
class robogauge:
enabled = False
port = 9973
class LeggedRobotCfgMoENGCTS(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
@@ -359,8 +367,8 @@ class LeggedRobotCfgMoECTS(LeggedRobotCfgCTS):
load_balance_coef = 0.01 # coefficient for load balance loss
class runner(LeggedRobotCfgCTS.runner):
policy_class_name = 'ActorCriticMoECTS'
algorithm_class_name = 'MoECTS'
policy_class_name = 'ActorCriticMoENGCTS'
algorithm_class_name = 'MoENGCTS'
class LeggedRobotCfgMCPCTS(LeggedRobotCfgCTS):
class policy(LeggedRobotCfgCTS.policy):
@@ -382,20 +390,20 @@ class LeggedRobotCfgACMoECTS(LeggedRobotCfgCTS):
class LeggedRobotCfgDualMoECTS(LeggedRobotCfgCTS):
class policy(LeggedRobotCfgCTS.policy):
expert_num = 8 # number of experts in the student model
student_encoder_hidden_dims = [512, 256, 128]
student_encoder_hidden_dims = [512, 256, 256]
class runner(LeggedRobotCfgCTS.runner):
policy_class_name = 'ActorCriticDualMoECTS'
algorithm_class_name = 'DualMoECTS'
class LeggedRobotCfgREMCTS(LeggedRobotCfgCTS):
class LeggedRobotCfgMoECTS(LeggedRobotCfgCTS):
class policy(LeggedRobotCfgCTS.policy):
expert_num = 8 # number of experts in the student model
student_encoder_hidden_dims = [512, 256, 128]
student_encoder_hidden_dims = [512, 256, 256]
class algorithm(LeggedRobotCfgCTS.algorithm):
load_balance_coef = 0.01 # coefficient for load balance loss
class runner(LeggedRobotCfgCTS.runner):
policy_class_name = 'ActorCriticREMCTS'
algorithm_class_name = 'REMCTS'
policy_class_name = 'ActorCriticMoECTS'
algorithm_class_name = 'MoECTS'

20
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, LeggedRobotCfgDualMoECTS, LeggedRobotCfgREMCTS
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoENGCTS, LeggedRobotCfgMoENGCTS, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS, LeggedRobotCfgDualMoECTS, LeggedRobotCfgMoECTS
class GO2Cfg(LeggedRobotCfg):
class init_state(LeggedRobotCfg.init_state):
@@ -228,17 +228,17 @@ class GO2CfgCTS(LeggedRobotCfgCTS):
latent_dim = 32
norm_type = 'l2norm'
class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
class policy(LeggedRobotCfgMoECTS.policy):
class GO2CfgMoENGCTS(LeggedRobotCfgMoENGCTS):
class policy(LeggedRobotCfgMoENGCTS.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 algorithm(LeggedRobotCfgMoECTS.algorithm):
class algorithm(LeggedRobotCfgMoENGCTS.algorithm):
load_balance_coef = 0.01
class runner(LeggedRobotCfgMoECTS.runner):
class runner(LeggedRobotCfgMoENGCTS.runner):
run_name = ''
experiment_name = 'go2_moe_cts'
experiment_name = 'go2_moe_no_goal_cts'
max_iterations = 150000
save_interval = 500
@@ -273,12 +273,12 @@ class GO2CfgDualMoECTS(LeggedRobotCfgDualMoECTS):
max_iterations = 150000
save_interval = 500
class GO2CfgREMCTS(LeggedRobotCfgREMCTS):
class policy(LeggedRobotCfgREMCTS.policy):
class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
class policy(LeggedRobotCfgMoECTS.policy):
expert_num = 8 # number of experts in the student model
class runner(LeggedRobotCfgREMCTS.runner):
class runner(LeggedRobotCfgMoECTS.runner):
run_name = ''
experiment_name = 'go2_rem_cts'
experiment_name = 'go2_moe_cts'
max_iterations = 150000
save_interval = 500

20
legged_gym/envs/go2/go2_config_fast_flat_move.py
View File

@@ -9,7 +9,7 @@
Change command_range_curriculum, init command range
'''
import math
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS, LeggedRobotCfgDualMoECTS, LeggedRobotCfgREMCTS
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoENGCTS, LeggedRobotCfgMoENGCTS, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS, LeggedRobotCfgDualMoECTS, LeggedRobotCfgMoECTS
class GO2Cfg(LeggedRobotCfg):
class init_state(LeggedRobotCfg.init_state):
@@ -259,17 +259,17 @@ class GO2CfgCTS(LeggedRobotCfgCTS):
latent_dim = 32
norm_type = 'l2norm'
class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
class policy(LeggedRobotCfgMoECTS.policy):
class GO2CfgMoENGCTS(LeggedRobotCfgMoENGCTS):
class policy(LeggedRobotCfgMoENGCTS.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 algorithm(LeggedRobotCfgMoECTS.algorithm):
class algorithm(LeggedRobotCfgMoENGCTS.algorithm):
load_balance_coef = 0.01
class runner(LeggedRobotCfgMoECTS.runner):
class runner(LeggedRobotCfgMoENGCTS.runner):
run_name = ''
experiment_name = 'go2_moe_cts'
experiment_name = 'go2_moe_no_goal_cts'
max_iterations = 150000
save_interval = 500
@@ -304,12 +304,12 @@ class GO2CfgDualMoECTS(LeggedRobotCfgDualMoECTS):
max_iterations = 150000
save_interval = 500
class GO2CfgREMCTS(LeggedRobotCfgREMCTS):
class policy(LeggedRobotCfgREMCTS.policy):
class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
class policy(LeggedRobotCfgMoECTS.policy):
expert_num = 8 # number of experts in the student model
class runner(LeggedRobotCfgREMCTS.runner):
class runner(LeggedRobotCfgMoECTS.runner):
run_name = ''
experiment_name = 'go2_rem_cts'
experiment_name = 'go2_moe_cts'
max_iterations = 150000
save_interval = 500

40
legged_gym/envs/go2/go2_config_vanilla.py
View File

@@ -1,17 +1,17 @@
# -*- coding: utf-8 -*-
'''
@File : go2_config_vanilla.py
@Time : 2026/01/10 02:26:04
@Time : 2026/01/10 02:27:28
@Author : wty-yy
@Version : 1.0
@Blog : https://wty-yy.github.io/
@Desc : Go2 vanilla training config
episode length 25, resample commands 5 sec,
open move_down_by_accumulated_xy_command, dynamic_resample_commands
close heading_command, zero_command_curriculum, limit_vel_prob, command_range_curriculum, dynamic_sigma
@Desc : Go2 vanilla training config, same as unitree rl gym except domain randomization and rewards
episode length 20, resample commands 10 sec,
open heading_command
close move_down_by_accumulated_xy_command, dynamic_resample_commands, zero_command_curriculum, limit_vel_prob, command_range_curriculum, dynamic_sigma
'''
import math
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS, LeggedRobotCfgDualMoECTS, LeggedRobotCfgREMCTS
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoENGCTS, LeggedRobotCfgMoENGCTS, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS, LeggedRobotCfgDualMoECTS, LeggedRobotCfgMoECTS
class GO2Cfg(LeggedRobotCfg):
class init_state(LeggedRobotCfg.init_state):
@@ -48,7 +48,7 @@ class GO2Cfg(LeggedRobotCfg):
num_privileged_obs = 45 + 3 + 4 + 12 + 12 + 187 # 263
# num_privileged_obs = 45 + 3 + 187 # 235
# num_privileged_obs = 48 # without height measurements
episode_length_s = 25
episode_length_s = 20
class domain_rand(LeggedRobotCfg.domain_rand):
### Robot properties ###
@@ -105,14 +105,14 @@ class GO2Cfg(LeggedRobotCfg):
# terrain_proportions = [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]
# terrain_proportions = [0.3, 0.3, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.1]
# terrain_proportions = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]
move_down_by_accumulated_xy_command = True # move down the terrain curriculum based on accumulated xy command distance instead of absolute distance
move_down_by_accumulated_xy_command = False # move down the terrain curriculum based on accumulated xy command distance instead of absolute distance
class commands(LeggedRobotCfg.commands):
curriculum = False
max_curriculum = 1.
num_commands = 4 # default: lin_vel_x, lin_vel_y, ang_vel_yaw (in heading mode ang_vel_yaw is recomputed from heading error)
resampling_time = 5. # time before command are changed[s]
heading_command = False # if true: compute ang vel command from heading error
resampling_time = 10. # time before command are changed[s]
heading_command = True # if true: compute ang vel command from heading error
# start training with zero commands and then gradually increase zero command probability
zero_command_curriculum = None
# zero_command_curriculum = {'start_iter': 0, 'end_iter': 1500, 'start_value': 0.0, 'end_value': 0.1}
@@ -121,7 +121,7 @@ class GO2Cfg(LeggedRobotCfg):
limit_vel_invert_when_continuous = True # invert the limit logic when using continuous sample limit velocity commands
limit_vel = {"lin_vel_x": [-1, 1], "lin_vel_y": [-1, 1], "ang_vel_yaw": [-1, 0, 1]} # sample vel commands from min [-1] or zero [0] or max [1] range only
stop_heading_at_limit = True # stop heading updates when vel is limited
dynamic_resample_commands = True # sample commands with low bounds
dynamic_resample_commands = False # sample commands with low bounds
command_range_curriculum = []
# command_range_curriculum = [{ # list for command range curriculums at specific training iterations
# 'iter': 20000, # training iteration at which the command ranges are updated
@@ -243,17 +243,17 @@ class GO2CfgCTS(LeggedRobotCfgCTS):
latent_dim = 32
norm_type = 'l2norm'
class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
class policy(LeggedRobotCfgMoECTS.policy):
class GO2CfgMoENGCTS(LeggedRobotCfgMoENGCTS):
class policy(LeggedRobotCfgMoENGCTS.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 algorithm(LeggedRobotCfgMoECTS.algorithm):
class algorithm(LeggedRobotCfgMoENGCTS.algorithm):
load_balance_coef = 0.01
class runner(LeggedRobotCfgMoECTS.runner):
class runner(LeggedRobotCfgMoENGCTS.runner):
run_name = ''
experiment_name = 'go2_moe_cts'
experiment_name = 'go2_moe_no_goal_cts'
max_iterations = 150000
save_interval = 500
@@ -288,12 +288,12 @@ class GO2CfgDualMoECTS(LeggedRobotCfgDualMoECTS):
max_iterations = 150000
save_interval = 500
class GO2CfgREMCTS(LeggedRobotCfgREMCTS):
class policy(LeggedRobotCfgREMCTS.policy):
class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
class policy(LeggedRobotCfgMoECTS.policy):
expert_num = 8 # number of experts in the student model
class runner(LeggedRobotCfgREMCTS.runner):
class runner(LeggedRobotCfgMoECTS.runner):
run_name = ''
experiment_name = 'go2_rem_cts'
experiment_name = 'go2_moe_cts'
max_iterations = 150000
save_interval = 500

42
legged_gym/envs/go2/go2_config_vanilla2.py → legged_gym/envs/go2/go2_config_vanilla_with_dynamic_cmd.py
View File

@@ -1,17 +1,17 @@
# -*- coding: utf-8 -*-
'''
@File : go2_config_vanilla2.py
@Time : 2026/01/10 02:27:28
@File : go2_config_vanilla_with_dynamic_cmd.py
@Time : 2026/01/10 02:26:04
@Author : wty-yy
@Version : 1.0
@Blog : https://wty-yy.github.io/
@Desc : Go2 vanilla2 training config, same as unitree rl gym except domain randomization and rewards
episode length 20, resample commands 10 sec,
open heading_command
close move_down_by_accumulated_xy_command, dynamic_resample_commands, zero_command_curriculum, limit_vel_prob, command_range_curriculum, dynamic_sigma
@Desc : Go2 vanilla training config
episode length 25, resample commands 5 sec,
open move_down_by_accumulated_xy_command, dynamic_resample_commands
close heading_command, zero_command_curriculum, limit_vel_prob, command_range_curriculum, dynamic_sigma
'''
import math
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS, LeggedRobotCfgDualMoECTS, LeggedRobotCfgREMCTS
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoENGCTS, LeggedRobotCfgMoENGCTS, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS, LeggedRobotCfgDualMoECTS, LeggedRobotCfgMoECTS
class GO2Cfg(LeggedRobotCfg):
class init_state(LeggedRobotCfg.init_state):
@@ -48,7 +48,7 @@ class GO2Cfg(LeggedRobotCfg):
num_privileged_obs = 45 + 3 + 4 + 12 + 12 + 187 # 263
# num_privileged_obs = 45 + 3 + 187 # 235
# num_privileged_obs = 48 # without height measurements
episode_length_s = 20
episode_length_s = 25
class domain_rand(LeggedRobotCfg.domain_rand):
### Robot properties ###
@@ -105,14 +105,14 @@ class GO2Cfg(LeggedRobotCfg):
# terrain_proportions = [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]
# terrain_proportions = [0.3, 0.3, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.1]
# terrain_proportions = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]
move_down_by_accumulated_xy_command = False # move down the terrain curriculum based on accumulated xy command distance instead of absolute distance
move_down_by_accumulated_xy_command = True # move down the terrain curriculum based on accumulated xy command distance instead of absolute distance
class commands(LeggedRobotCfg.commands):
curriculum = False
max_curriculum = 1.
num_commands = 4 # default: lin_vel_x, lin_vel_y, ang_vel_yaw (in heading mode ang_vel_yaw is recomputed from heading error)
resampling_time = 10. # time before command are changed[s]
heading_command = True # if true: compute ang vel command from heading error
resampling_time = 5. # time before command are changed[s]
heading_command = False # if true: compute ang vel command from heading error
# start training with zero commands and then gradually increase zero command probability
zero_command_curriculum = None
# zero_command_curriculum = {'start_iter': 0, 'end_iter': 1500, 'start_value': 0.0, 'end_value': 0.1}
@@ -121,7 +121,7 @@ class GO2Cfg(LeggedRobotCfg):
limit_vel_invert_when_continuous = True # invert the limit logic when using continuous sample limit velocity commands
limit_vel = {"lin_vel_x": [-1, 1], "lin_vel_y": [-1, 1], "ang_vel_yaw": [-1, 0, 1]} # sample vel commands from min [-1] or zero [0] or max [1] range only
stop_heading_at_limit = True # stop heading updates when vel is limited
dynamic_resample_commands = False # sample commands with low bounds
dynamic_resample_commands = True # sample commands with low bounds
command_range_curriculum = []
# command_range_curriculum = [{ # list for command range curriculums at specific training iterations
# 'iter': 20000, # training iteration at which the command ranges are updated
@@ -243,17 +243,17 @@ class GO2CfgCTS(LeggedRobotCfgCTS):
latent_dim = 32
norm_type = 'l2norm'
class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
class policy(LeggedRobotCfgMoECTS.policy):
class GO2CfgMoENGCTS(LeggedRobotCfgMoENGCTS):
class policy(LeggedRobotCfgMoENGCTS.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 algorithm(LeggedRobotCfgMoECTS.algorithm):
class algorithm(LeggedRobotCfgMoENGCTS.algorithm):
load_balance_coef = 0.01
class runner(LeggedRobotCfgMoECTS.runner):
class runner(LeggedRobotCfgMoENGCTS.runner):
run_name = ''
experiment_name = 'go2_moe_cts'
experiment_name = 'go2_moe_no_goal_cts'
max_iterations = 150000
save_interval = 500
@@ -288,12 +288,12 @@ class GO2CfgDualMoECTS(LeggedRobotCfgDualMoECTS):
max_iterations = 150000
save_interval = 500
class GO2CfgREMCTS(LeggedRobotCfgREMCTS):
class policy(LeggedRobotCfgREMCTS.policy):
class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
class policy(LeggedRobotCfgMoECTS.policy):
expert_num = 8 # number of experts in the student model
class runner(LeggedRobotCfgREMCTS.runner):
class runner(LeggedRobotCfgMoECTS.runner):
run_name = ''
experiment_name = 'go2_rem_cts'
experiment_name = 'go2_moe_cts'
max_iterations = 150000
save_interval = 500

18
legged_gym/utils/exporter.py
View File

@@ -81,9 +81,9 @@ class _TorchPolicyExporter(torch.nn.Module):
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
self.forward = self.forward_moe_no_goal_cts
if not hasattr(policy, "obs_no_goal_mask"):
self.forward = self.forward_rem_cts
self.forward = self.forward_moe_cts
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()
@@ -134,7 +134,7 @@ class _TorchPolicyExporter(torch.nn.Module):
x = torch.cat([latent, x], dim=1)
return self.actor(x), (None, latent)
def forward_moe_cts(self, x): # x is single observations
def forward_moe_no_goal_cts(self, x): # x is single observations
x = self.normalizer(x)
self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
history_no_goal = self.history.reshape(1, self.history_length, -1)[:, :, self.obs_no_goal_mask].reshape(1, -1)
@@ -142,7 +142,7 @@ class _TorchPolicyExporter(torch.nn.Module):
x = torch.cat([latent, x], dim=1)
return self.actor(x), (weights, latent)
def forward_rem_cts(self, x): # x is single observations
def forward_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, weights = self.student_moe_encoder(self.history.flatten(1))
@@ -211,12 +211,12 @@ class _OnnxPolicyExporter(torch.nn.Module):
elif hasattr(policy, "student_moe_encoder"):
self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
self.history_length = policy.history.shape[1]
self.forward = self.forward_moe_cts
self.forward = self.forward_moe_no_goal_cts
self.input_dim = self.history_length * policy.history.shape[2]
if hasattr(policy, "obs_no_goal_mask"):
self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
else:
self.forward = self.forward_rem_cts
self.forward = self.forward_moe_cts
else: # PPO
self.forward = self.forward_ppo
@@ -277,7 +277,7 @@ class _OnnxPolicyExporter(torch.nn.Module):
return self.actor(x)
def forward_moe_cts(self, x):
def forward_moe_no_goal_cts(self, x):
x = self.normalizer(x)
history, obs_dim = self.flatten_obs(x)
@@ -290,7 +290,7 @@ class _OnnxPolicyExporter(torch.nn.Module):
return self.actor(x), weights, latent
def forward_rem_cts(self, x):
def forward_moe_cts(self, x):
x = self.normalizer(x)
history, obs_dim = self.flatten_obs(x)
@@ -319,7 +319,7 @@ class _OnnxPolicyExporter(torch.nn.Module):
obs = torch.zeros(1, self.input_dim)
output_names = ["actions"]
if self.forward == self.forward_moe_cts:
if self.forward == self.forward_moe_no_goal_cts:
output_names.append("weights")
output_names.append("latent")
if self.forward == self.forward_mcp_cts:

7
legged_gym/utils/helpers.py
View File

@@ -118,6 +118,10 @@ def update_cfg_from_args(env_cfg, cfg_train, args):
cfg_train.runner.load_run = args.load_run
if args.checkpoint is not None:
cfg_train.runner.checkpoint = args.checkpoint
if args.robogauge is not None:
cfg_train.robogauge.enabled = args.robogauge
if args.robogauge_port is not None:
cfg_train.robogauge.port = args.robogauge_port
return env_cfg, cfg_train
@@ -136,6 +140,9 @@ def get_args():
{"name": "--num_envs", "type": int, "help": "Number of environments to create. Overrides config file if provided."},
{"name": "--seed", "type": int, "help": "Random seed. Overrides config file if provided."},
{"name": "--max_iterations", "type": int, "help": "Maximum number of training iterations. Overrides config file if provided."},
{"name": "--robogauge", "action": "store_true", "default": False, "help": "Enable robogauge evaluation interface."},
{"name": "--robogauge_port", "type": int, "default": 9973, "help": "Port for robogauge evaluation interface."},
]
# parse arguments
args = gymutil.parse_arguments(