Init

legged_gym/envs/__init__.py

@@ -0,0 +1,11 @@
+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 .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())

legged_gym/envs/base/base_config.py

@@ -0,0 +1,25 @@
+import inspect
+
+class BaseConfig:
+    def __init__(self) -> None:
+        """ Initializes all member classes recursively. Ignores all namse starting with '__' (buit-in methods)."""
+        self.init_member_classes(self)
+    
+    @staticmethod
+    def init_member_classes(obj):
+        # iterate over all attributes names
+        for key in dir(obj):
+            # disregard builtin attributes
+            # if key.startswith("__"):
+            if key=="__class__":
+                continue
+            # get the corresponding attribute object
+            var =  getattr(obj, key)
+            # check if it the attribute is a class
+            if inspect.isclass(var):
+                # instantate the class
+                i_var = var()
+                # set the attribute to the instance instead of the type
+                setattr(obj, key, i_var)
+                # recursively init members of the attribute
+                BaseConfig.init_member_classes(i_var)

legged_gym/envs/base/base_task.py

@@ -0,0 +1,115 @@
+
+import sys
+from isaacgym import gymapi
+from isaacgym import gymutil
+import numpy as np
+import torch
+
+# Base class for RL tasks
+class BaseTask():
+
+    def __init__(self, cfg, sim_params, physics_engine, sim_device, headless):
+        self.gym = gymapi.acquire_gym()
+
+        self.sim_params = sim_params
+        self.physics_engine = physics_engine
+        self.sim_device = sim_device
+        sim_device_type, self.sim_device_id = gymutil.parse_device_str(self.sim_device)
+        self.headless = headless
+
+        # env device is GPU only if sim is on GPU and use_gpu_pipeline=True, otherwise returned tensors are copied to CPU by physX.
+        if sim_device_type=='cuda' and sim_params.use_gpu_pipeline:
+            self.device = self.sim_device
+        else:
+            self.device = 'cpu'
+
+        # graphics device for rendering, -1 for no rendering
+        self.graphics_device_id = self.sim_device_id
+        if self.headless == True:
+            self.graphics_device_id = -1
+
+        self.num_envs = cfg.env.num_envs
+        self.num_obs = cfg.env.num_observations
+        self.num_privileged_obs = cfg.env.num_privileged_obs
+        self.num_actions = cfg.env.num_actions
+
+        # optimization flags for pytorch JIT
+        torch._C._jit_set_profiling_mode(False)
+        torch._C._jit_set_profiling_executor(False)
+
+        # allocate buffers
+        self.obs_buf = torch.zeros(self.num_envs, self.num_obs, device=self.device, dtype=torch.float)
+        self.rew_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.float)
+        self.reset_buf = torch.ones(self.num_envs, device=self.device, dtype=torch.long)
+        self.episode_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long)
+        self.time_out_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool)
+        if self.num_privileged_obs is not None:
+            self.privileged_obs_buf = torch.zeros(self.num_envs, self.num_privileged_obs, device=self.device, dtype=torch.float)
+        else: 
+            self.privileged_obs_buf = None
+            # self.num_privileged_obs = self.num_obs
+
+        self.extras = {}
+
+        # create envs, sim and viewer
+        self.create_sim()
+        self.gym.prepare_sim(self.sim)
+
+        # todo: read from config
+        self.enable_viewer_sync = True
+        self.viewer = None
+
+        # if running with a viewer, set up keyboard shortcuts and camera
+        if self.headless == False:
+            # subscribe to keyboard shortcuts
+            self.viewer = self.gym.create_viewer(
+                self.sim, gymapi.CameraProperties())
+            self.gym.subscribe_viewer_keyboard_event(
+                self.viewer, gymapi.KEY_ESCAPE, "QUIT")
+            self.gym.subscribe_viewer_keyboard_event(
+                self.viewer, gymapi.KEY_V, "toggle_viewer_sync")
+
+    def get_observations(self):
+        return self.obs_buf
+    
+    def get_privileged_observations(self):
+        return self.privileged_obs_buf
+
+    def reset_idx(self, env_ids):
+        """Reset selected robots"""
+        raise NotImplementedError
+
+    def reset(self):
+        """ Reset all robots"""
+        self.reset_idx(torch.arange(self.num_envs, device=self.device))
+        obs, privileged_obs, _, _, _ = self.step(torch.zeros(self.num_envs, self.num_actions, device=self.device, requires_grad=False))
+        return obs, privileged_obs
+
+    def step(self, actions):
+        raise NotImplementedError
+
+    def render(self, sync_frame_time=True):
+        if self.viewer:
+            # check for window closed
+            if self.gym.query_viewer_has_closed(self.viewer):
+                sys.exit()
+
+            # check for keyboard events
+            for evt in self.gym.query_viewer_action_events(self.viewer):
+                if evt.action == "QUIT" and evt.value > 0:
+                    sys.exit()
+                elif evt.action == "toggle_viewer_sync" and evt.value > 0:
+                    self.enable_viewer_sync = not self.enable_viewer_sync
+
+            # fetch results
+            if self.device != 'cpu':
+                self.gym.fetch_results(self.sim, True)
+
+            # step graphics
+            if self.enable_viewer_sync:
+                self.gym.step_graphics(self.sim)
+                self.gym.draw_viewer(self.viewer, self.sim, True)
+                if sync_frame_time:
+                    self.gym.sync_frame_time(self.sim)
+            else:
+                self.gym.poll_viewer_events(self.viewer)

legged_gym/envs/base/legged_robot_config.py

@@ -0,0 +1,361 @@
+import math
+from .base_config import BaseConfig
+
+class LeggedRobotCfg(BaseConfig):
+    class env:
+        num_envs = 4096
+        num_observations = 48
+        num_privileged_obs = None # if not None a priviledge_obs_buf will be returned by step() (critic obs for assymetric training). None is returned otherwise 
+        num_actions = 12
+        env_spacing = 3.  # not used with heightfields/trimeshes 
+        send_timeouts = True # send time out information to the algorithm
+        episode_length_s = 20 # episode length in seconds
+        test = False
+
+    class terrain:
+        mesh_type = 'trimesh' # none, plane, heightfield or trimesh
+        horizontal_scale = 0.1 # [m]
+        vertical_scale = 0.005 # [m]
+        border_size = 25 # [m]
+        curriculum = True
+        static_friction = 1.0
+        dynamic_friction = 1.0
+        restitution = 0.
+        # rough terrain only:
+        measure_heights = True
+        measured_points_x = [-0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1, 0., 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8] # 1mx1.6m rectangle (without center line)
+        measured_points_y = [-0.5, -0.4, -0.3, -0.2, -0.1, 0., 0.1, 0.2, 0.3, 0.4, 0.5]
+        selected = False # select a unique terrain type and pass all arguments
+        terrain_kwargs = None # Dict of arguments for selected terrain
+        max_init_terrain_level = 5 # starting curriculum state
+        terrain_length = 8.
+        terrain_width = 8.
+        num_rows= 10 # number of terrain rows (levels)
+        num_cols = 20 # number of terrain cols (types)
+        terrain_spacing = 0.5 # spacing between different terrain types [m]
+
+        # [wave, slope, rough slope, stairs down, stairs up, obstacles, stepping stones, gap, flat]
+        terrain_proportions = [0.1, 0.1, 0.1, 0.2, 0.2, 0.1, 0.1, 0.1, 0.0]
+        # trimesh only:
+        slope_treshold = 0.75 # slopes above this threshold will be corrected to vertical surfaces
+
+    class commands:
+        curriculum = False
+        max_curriculum = 1.
+        num_commands = 4 # default: lin_vel_x, lin_vel_y, ang_vel_yaw, heading (in heading mode ang_vel_yaw is recomputed from heading error)
+        resampling_time = 10. # time before command are changed[s]
+        heading_command = False # if true: compute ang vel command from heading error
+        zero_command_curriculum = None
+        # start training with zero commands and then gradually increase zero command probability
+        # eg. {'start_iter': 0, 'end_iter': 1500, 'start_value': 0.0, 'end_value': 0.1}
+        limit_ang_vel_at_zero_command_prob = 0.0 # probability of add limiting angular velocity commands when zero command is sampled
+        limit_vel_prob = 0.0 # probability of limiting linear velocity command
+        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
+        command_range_curriculum = [] # list for command range curriculums at specific training iterations
+        # eg: [{
+        #     'iter': 20000, # training iteration at which the command ranges are updated
+        #     'lin_vel_x': [-1.0, 1.0], # min max [m/s]
+        #     'lin_vel_y': [-1.0, 1.0], # min max [m/s]
+        #     'ang_vel_yaw': [-2.0, 2.0], # min max [rad/s]
+        #     'heading': [-1.57, 1.57], # min max [rad]
+        # }]
+        turn_over_zero_time = { # if turn_over is true, time robot must be stable before sampling new commands after a turn over
+            "backflip": 5.0,
+            "sideflip": 3.0,
+        }
+        # [wave, slope, rough slope, stairs up, stairs down, obstacles, stepping stones, gap, flat]
+        terrain_max_command_ranges = [
+            {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.5, 1.5], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # wave
+            {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.5, 1.5], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # slope
+            {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.5, 1.5], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # rough slope
+            {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # stairs up
+            {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # stairs down
+            {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # obstacles
+            {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # stepping stones
+            {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # gap
+            {'lin_vel_x': [-2.0, 2.0], 'lin_vel_y': [-1.5, 1.5], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # flat
+        ]
+
+        class ranges:
+            lin_vel_x = [-1.0, 1.0] # min max [m/s]
+            lin_vel_y = [-0.5, 0.5]   # min max [m/s]
+            ang_vel_yaw = [-1, 1]    # min max [rad/s]
+            heading = [-3.14, 3.14]
+
+    class init_state:
+        pos = [0.0, 0.0, 1.] # x,y,z [m]
+        rot = [0.0, 0.0, 0.0, 1.0] # x,y,z,w [quat]
+        lin_vel = [0.0, 0.0, 0.0]  # x,y,z [m/s]
+        ang_vel = [0.0, 0.0, 0.0]  # x,y,z [rad/s]
+        default_joint_angles = { # target angles when action = 0.0
+            "joint_a": 0., 
+            "joint_b": 0.}
+        turn_over = False # if true, initialize the robot in a flipped over position
+        turn_over_proportions = [0.0, 0.2, 0.8] # proportions for backflip, sideflip, no flip
+        turn_over_init_heights = { # initial heights range for each flip type
+            'backflip': [0.10, 0.15],
+            'sideflip': [0.16, 0.21],
+        }
+
+    class control:
+        control_type = 'P' # P: position, V: velocity, T: torques
+        # PD Drive parameters:
+        stiffness = {'joint_a': 10.0, 'joint_b': 15.}  # [N*m/rad]
+        damping = {'joint_a': 1.0, 'joint_b': 1.5}     # [N*m*s/rad]
+        # action scale: target angle = actionScale * action + defaultAngle
+        action_scale = 0.5
+        # decimation: Number of control action updates @ sim DT per policy DT
+        decimation = 4
+
+    class asset:
+        file = ""
+        name = "legged_robot"  # actor name
+        foot_name = "None" # name of the feet bodies, used to index body state and contact force tensors
+        penalize_contacts_on = []
+        terminate_after_contacts_on = []
+        disable_gravity = False
+        collapse_fixed_joints = True # merge bodies connected by fixed joints. Specific fixed joints can be kept by adding " <... dont_collapse="true">
+        fix_base_link = False # fixe the base of the robot
+        default_dof_drive_mode = 3 # see GymDofDriveModeFlags (0 is none, 1 is pos tgt, 2 is vel tgt, 3 effort)
+        self_collisions = 0 # 1 to disable, 0 to enable...bitwise filter
+        replace_cylinder_with_capsule = True # replace collision cylinders with capsules, leads to faster/more stable simulation
+        flip_visual_attachments = True # Some .obj meshes must be flipped from y-up to z-up
+        
+        density = 0.001
+        angular_damping = 0.
+        linear_damping = 0.
+        max_angular_velocity = 1000.
+        max_linear_velocity = 1000.
+        armature = 0.
+        thickness = 0.01
+
+    class domain_rand:
+        ### Robot properties ###
+        robot_properties_update = None
+        # eg: {'start_iter': 5000, 'interval': 5000}
+
+        randomize_friction = True
+        friction_range = [0.2, 1.25]
+
+        randomize_base_mass = True
+        added_mass_range = [-1., 1.]
+
+        randomize_link_mass = True
+        multiplied_link_mass_range = [0.9, 1.1]
+
+        randomize_base_com = True
+        added_base_com_range = [-0.03, 0.03]
+
+        randomize_restitution = False # restitution to robot links (Robot init)
+        restitution_range = [0.0, 0.2]
+
+        ### Environment reset ###
+        randomize_pd_gains = True
+        stiffness_multiplier_range = [0.9, 1.1]  
+        damping_multiplier_range = [0.9, 1.1]    
+
+        randomize_motor_zero_offset = True
+        motor_zero_offset_range = [-0.035, 0.035]
+
+        randomize_motor_strength = False # (Env reset)
+        motor_strength_range = [0.8, 1.2]
+
+        ### Environment step ###
+        push_robots = True
+        push_interval_s = 4
+        max_push_vel_xy = 0.4
+        max_push_ang_vel = 0.6
+
+        randomize_action_delay = False # use last_action with 0~20 ms delay, 4 decimation
+
+    class rewards:
+        class scales:
+            termination = -0.0
+            tracking_lin_vel = 1.0
+            tracking_ang_vel = 0.5
+            lin_vel_z = -2.0
+            ang_vel_xy = -0.05
+            orientation = -0.
+            torques = -0.00001
+            dof_vel = -0.
+            dof_acc = -2.5e-7
+            base_height = -0. 
+            feet_air_time = 1.0
+            collision = -1.
+            feet_stumble = -0.0 
+            action_rate = -0.01
+            stand_still = -0.
+
+        class turn_over_scales:
+            upright = 1.0
+
+        only_positive_rewards = True # if true negative total rewards are clipped at zero (avoids early termination problems)
+        tracking_sigma = 0.25 # tracking reward = exp(-error^2/sigma)
+        soft_dof_pos_limit = 1. # percentage of urdf limits, values above this limit are penalized
+        soft_dof_vel_limit = 1.
+        soft_torque_limit = 1.
+        base_height_target = 1.
+        max_contact_force = 100. # forces above this value are penalized
+        curriculum_rewards = None  # reward names to apply curriculum scaling to, List[dict]
+        # eg: [{'reward_name': 'lin_vel_z', 'start_iter': 0, 'end_iter': 1500, 'start_value': 1.0, 'end_value': 0.0}]
+        dynamic_sigma = None # linear interpolation of sigma based on command velocity, **Must start terrain curriculum first**
+        # eg: {
+        #     "min_vel": 0.5, # min abs velocity to have default sigma
+        #     "max_vel": 1.0, # max abs velocity to have max sigma
+        #     # wave, slope, rough_slope, stairs up, stairs down, obstacles, stepping_stones, gap, flat]
+        #     "max_sigma": [1/3, 1/4, 1/4, 1/2.7, 1/2.7, 1/2, 1, 1, 1/4]
+        # }
+        turn_over_roll_threshold = math.pi / 4 # threshold on roll to use turn over rewards
+        min_legs_distance = 0.1  # min distance between legs to not be considered stumbling
+
+    class normalization:
+        class obs_scales:
+            lin_vel = 2.0
+            ang_vel = 0.25
+            dof_pos = 1.0
+            dof_vel = 0.05
+            height_measurements = 2.5
+        clip_observations = 100.
+        clip_actions = 100.
+
+    class noise:
+        add_noise = True
+        noise_level = 1.0 # scales other values
+        class noise_scales:
+            dof_pos = 0.01
+            dof_vel = 1.5
+            lin_vel = 0.1
+            ang_vel = 0.2
+            gravity = 0.05
+            height_measurements = 0.1
+
+    # viewer camera:
+    class viewer:
+        ref_env = 0
+        pos = [10, 0, 6]  # [m]
+        lookat = [11., 5, 3.]  # [m]
+
+    class sim:
+        dt =  0.005
+        substeps = 1
+        gravity = [0., 0. ,-9.81]  # [m/s^2]
+        up_axis = 1  # 0 is y, 1 is z
+
+        class physx:
+            num_threads = 10
+            solver_type = 1  # 0: pgs, 1: tgs
+            num_position_iterations = 4
+            num_velocity_iterations = 0
+            contact_offset = 0.01  # [m]
+            rest_offset = 0.0   # [m]
+            bounce_threshold_velocity = 0.5 #0.5 [m/s]
+            max_depenetration_velocity = 1.0
+            max_gpu_contact_pairs = 2**23 #2**24 -> needed for 8000 envs and more
+            default_buffer_size_multiplier = 5
+            contact_collection = 2 # 0: never, 1: last sub-step, 2: all sub-steps (default=2)
+
+class LeggedRobotCfgPPO(BaseConfig):
+    seed = 1
+    runner_class_name = 'OnPolicyRunner'
+    class policy:
+        init_noise_std = 1.0
+        actor_hidden_dims = [512, 256, 128]
+        critic_hidden_dims = [512, 256, 128]
+        activation = 'elu' # can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
+        # only for 'ActorCriticRecurrent':
+        # rnn_type = 'lstm'
+        # rnn_hidden_size = 512
+        # rnn_num_layers = 1
+        
+    class algorithm:
+        # training params
+        value_loss_coef = 1.0
+        use_clipped_value_loss = True
+        clip_param = 0.2
+        entropy_coef = 0.01
+        num_learning_epochs = 5
+        num_mini_batches = 4 # mini batch size = num_envs*nsteps / nminibatches
+        learning_rate = 1.e-3 #5.e-4
+        schedule = 'adaptive' # could be adaptive, fixed
+        gamma = 0.99
+        lam = 0.95
+        desired_kl = 0.01
+        max_grad_norm = 1.
+
+    class runner:
+        policy_class_name = 'ActorCritic'
+        algorithm_class_name = 'PPO'
+        num_steps_per_env = 24 # per iteration
+        max_iterations = 1500 # number of policy updates
+
+        # logging
+        save_interval = 50 # check for potential saves every this many iterations
+        experiment_name = 'test'
+        run_name = ''
+        # load and resume
+        resume = False
+        load_run = -1 # -1 = last run
+        checkpoint = -1 # -1 = last saved model
+        resume_path = None # updated from load_run and chkpt
+
+class LeggedRobotCfgCTS(BaseConfig):
+    seed = 0
+    runner_class_name = "OnPolicyRunnerCTS"
+    history_length = 5
+    class policy:
+        init_noise_std = 1.0
+        actor_hidden_dims = [512, 256, 128]
+        critic_hidden_dims = [512, 256, 128]
+        teacher_encoder_hidden_dims = [512, 256]
+        student_encoder_hidden_dims = [512, 256]
+        activation = 'elu' # can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
+        latent_dim = 32
+        norm_type = 'l2norm' # normalization type for encoders: l2norm, simnorm
+    
+    class algorithm:
+        # training params
+        value_loss_coef = 1.0
+        use_clipped_value_loss = True
+        clip_param = 0.2
+        entropy_coef = 0.01
+        num_learning_epochs = 5
+        num_mini_batches = 4 # mini batch size = num_envs*nsteps / nminibatches
+        learning_rate = 1.e-3 #5.e-4
+        student_encoder_learning_rate = 1e-3
+        schedule = 'adaptive' # could be adaptive, fixed
+        gamma = 0.99
+        lam = 0.95
+        desired_kl = 0.01
+        max_grad_norm = 1.
+        teacher_env_ratio = 0.75  # percentage of envs assigned to teacher
+        # teacher_env_ratio = 1.00  # percentage of envs assigned to teacher
+
+    class runner:
+        policy_class_name = 'ActorCriticCTS'
+        algorithm_class_name = 'CTS'
+        num_steps_per_env = 24 # per iteration
+        max_iterations = 1500 # number of policy updates
+
+        # logging
+        save_interval = 50  # check for potential saves every this many iterations
+        experiment_name = 'test'
+        run_name = ''
+        # load and resume
+        resume = False
+        load_run = -1 # -1 = last run
+        checkpoint = -1 # -1 = last saved model
+        resume_path = None # updated from load_run and chkpt
+
+class LeggedRobotCfgMoECTS(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 algorithm(LeggedRobotCfgCTS.algorithm):
+        load_balance_coef = 0.01  # coefficient for load balance loss
+
+    class runner(LeggedRobotCfgCTS.runner):
+        policy_class_name = 'ActorCriticMoECTS'
+        algorithm_class_name = 'MoECTS'

legged_gym/envs/go2/go2_config.py

@@ -0,0 +1,266 @@
+import math
+from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS
+
+class GO2Cfg(LeggedRobotCfg):
+    class init_state(LeggedRobotCfg.init_state):
+        pos = [0.0, 0.0, 0.42] # x,y,z [m]
+        default_joint_angles = { # = target angles [rad] when action = 0.0
+            'FL_hip_joint': 0.1,   # [rad]
+            'RL_hip_joint': 0.1,   # [rad]
+            'FR_hip_joint': -0.1 ,  # [rad]
+            'RR_hip_joint': -0.1,   # [rad]
+
+            'FL_thigh_joint': 0.8,     # [rad]
+            'RL_thigh_joint': 1.,   # [rad]
+            'FR_thigh_joint': 0.8,     # [rad]
+            'RR_thigh_joint': 1.,   # [rad]
+
+            'FL_calf_joint': -1.5,   # [rad]
+            'RL_calf_joint': -1.5,    # [rad]
+            'FR_calf_joint': -1.5,  # [rad]
+            'RR_calf_joint': -1.5,    # [rad]
+        }
+        turn_over = False # initialize the robot in a flipped over position
+        # turn_over_proportions = [0.1, 0.3, 0.6] # proportions for backflip, sideflip, noflip
+        turn_over_proportions = [0.0, 0.2, 0.8] # proportions for backflip, sideflip, noflip
+        turn_over_init_heights = { # initial heights range for each flip type
+            'backflip': [0.10, 0.15],
+            'sideflip': [0.16, 0.21],
+        }
+        # turn_over_proportions = [0.0, 1.0, 0.0] # proportions for backflip, sideflip, noflip
+
+    class env(LeggedRobotCfg.env):
+        num_envs = 8192
+        num_observations = 45
+        # obs(45) + base_lin_vel(3) + height_measurements(187)
+        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
+
+    class domain_rand(LeggedRobotCfg.domain_rand):
+        ### Robot properties ###
+        randomize_friction = True
+        friction_range = [0.0, 2.0]
+
+        randomize_base_mass = True
+        added_mass_range = [-1., 1.]
+
+        randomize_link_mass = True
+        multiplied_link_mass_range = [0.9, 1.1]
+
+        randomize_base_com = True
+        added_base_com_range = [-0.03, 0.03]
+
+        randomize_restitution = True # restitution to robot links (Robot init)
+        restitution_range = [0.0, 0.5]
+
+        ### Environment reset ###
+        randomize_pd_gains = True
+        stiffness_multiplier_range = [0.9, 1.1]  
+        damping_multiplier_range = [0.9, 1.1]    
+
+        randomize_motor_zero_offset = True
+        motor_zero_offset_range = [-0.035, 0.035]
+
+        randomize_motor_strength = True # (Env reset)
+        motor_strength_range = [0.8, 1.2]
+
+        ### Environment step ###
+        push_robots = True
+        push_interval_s = 4
+        max_push_vel_xy = 0.4
+        max_push_ang_vel = 0.6
+
+        randomize_action_delay = True # use last_action with 0~20 ms delay, 4 decimation
+
+    class control(LeggedRobotCfg.control):
+        # PD Drive parameters:
+        control_type = 'P'
+        stiffness = {'joint': 20.0}  # [N*m/rad]
+        damping = {'joint': 0.5}     # [N*m*s/rad]
+        # action scale: target angle = actionScale * action + defaultAngle
+        action_scale = 0.25
+        # decimation: Number of control action updates @ sim DT per policy DT
+        decimation = 4
+    
+    class terrain(LeggedRobotCfg.terrain):
+        max_init_terrain_level = 5
+        # [wave, slope, rough_slope, stairs up, stairs down, obstacles, stepping_stones, gap, flat]
+        # terrain_proportions = [0.2, 0.05, 0.05, 0.30, 0.05, 0.25, 0.0, 0.0, 0.1]  # 更偏向wave
+        terrain_proportions = [0.05, 0.20, 0.05, 0.25, 0.10, 0.20, 0.0, 0.0, 0.15]  # 这个更偏向平地斜坡
+        # terrain_proportions = [0.20, 0.05, 0.05, 0.30, 0.15, 0.20, 0.0, 0.0, 0.05]  # 更偏向wave和stairs
+        # 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]
+        
+    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
+        # start training with zero commands and then gradually increase zero command probability
+        zero_command_curriculum = {'start_iter': 0, 'end_iter': 1500, 'start_value': 0.0, 'end_value': 0.1}
+        limit_ang_vel_at_zero_command_prob = 0.2 # probability of add limiting angular velocity commands when zero command is sampled
+        limit_vel_prob = 0.2 # probability of limiting linear velocity command
+        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
+        command_range_curriculum = [{ # list for command range curriculums at specific training iterations
+            'iter': 20000, # training iteration at which the command ranges are updated
+            'lin_vel_x': [-1.0, 1.0], # min max [m/s]
+            'lin_vel_y': [-1.0, 1.0], # min max [m/s]
+            'ang_vel_yaw': [-1.5, 1.5], # min max [rad/s]
+            'heading': [-1.57, 1.57], # min max [rad]
+        }, { # list for command range curriculums at specific training iterations
+            'iter': 50000, # training iteration at which the command ranges are updated
+            'lin_vel_x': [-2.0, 2.0], # min max [m/s]
+            'lin_vel_y': [-1.0, 1.0], # min max [m/s]
+            'ang_vel_yaw': [-2.0, 2.0], # min max [rad/s]
+            'heading': [-1.57, 1.57], # min max [rad]
+        }]
+        turn_over_zero_time = { # if turn_over is true, time robot must be stable before sampling new commands after a turn over
+            "backflip": 5.0,
+            "sideflip": 3.0,
+        }
+        # [wave, slope, rough slope, stairs up, stairs down, obstacles, stepping stones, gap, flat]
+        terrain_max_command_ranges = [
+            {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # wave
+            {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # slope
+            {'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # rough slope
+            {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # stairs up
+            {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # stairs down
+            {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # obstacles
+            {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # stepping stones
+            {'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},  # gap
+            {'lin_vel_x': [-2.0, 2.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-2.0, 2.0], 'heading': [-1.57, 1.57]},  # flat
+        ]
+
+        class ranges:
+            lin_vel_x = [-0.5, 0.5] # min max [m/s]
+            lin_vel_y = [-0.5, 0.5] # min max [m/s]
+            ang_vel_yaw = [-1.0, 1.0]   # min max [rad/s]
+            heading = [-1.57, 1.57] # min max [rad]
+        
+    class asset(LeggedRobotCfg.asset):
+        file = '{LEGGED_GYM_ROOT_DIR}/resources/robots/go2/urdf/go2.urdf'
+        name = "go2"
+        foot_name = "foot"
+        penalize_contacts_on = ["thigh", "calf"]
+        terminate_after_contacts_on = ["base"]
+        self_collisions = 1 # 1 to disable, 0 to enable...bitwise filter
+  
+    class rewards(LeggedRobotCfg.rewards):
+        soft_dof_pos_limit = 0.9
+        base_height_target = 0.38
+        only_positive_rewards = False
+        max_contact_force = 147. # forces above this value are penalized, go2 weight 15kg
+        curriculum_rewards = [
+            {'reward_name': 'lin_vel_z', 'start_iter': 0, 'end_iter': 1500, 'start_value': 1.0, 'end_value': 0.0},
+            {'reward_name': 'correct_base_height', 'start_iter': 0, 'end_iter': 5000, 'start_value': 1.0, 'end_value': 10.0},
+            # {'reward_name': 'dof_power', 'start_iter': 0, 'end_iter': 3000, 'start_value': 1.0, 'end_value': 0.1},
+            # {'reward_name': 'upright', 'start_iter': 0, 'end_iter': 1500, 'start_value': 1.0, 'end_value': 0.0},
+        ]
+        tracking_sigma = 0.25 # tracking reward = exp(-error^2/sigma)
+        dynamic_sigma = { # linear interpolation of sigma based on command velocity, **Must start terrain curriculum first**
+            "min_lin_vel": 0.5, # min abs linear velocity to have default sigma
+            "max_lin_vel": 1.5, # max abs linear velocity to have max sigma
+            "min_ang_vel": 1.0, # min abs angular velocity to have default sigma
+            "max_ang_vel": 2.0, # max abs angular velocity to have max sigma
+            # wave, slope, rough_slope, stairs up, stairs down, obstacles, stepping_stones, gap, flat]
+            # "max_sigma": [1/3, 1/4, 1/4, 1/2.7, 1/2.7, 1/2, 1, 1, 1/4]
+            "max_sigma": [5/12, 1/4, 1/4, 1/2, 1/2, 3/4, 1, 1, 1/4]
+        }
+        min_legs_distance = 0.1  # min distance between legs to not be considered stumbling
+        class scales:
+            # tracking_lin_vel = 1.0
+            # tracking_ang_vel = 0.2
+            # lin_vel_z = -10.0
+            # base_height = -50.0
+            # action_rate = -0.005
+            # similar_to_default = -0.1
+            # dof_power = -1e-3  # 能够明显抑制跳跃
+            # dof_acc = -3e-7
+
+            # tracking_lin_vel = 1.0
+            # tracking_ang_vel = 0.5
+            # lin_vel_z = -2.0
+            # ang_vel_xy = -0.05
+            # dof_acc = -2.5e-7
+            # dof_power = -1e-3  # 能够明显抑制跳跃
+            # # torques = -1e-4  # 无用会走着走着倒了
+            # correct_base_height = -10.0
+            # action_rate = -0.01
+            # action_smoothness = -0.01
+            # collision = -1.0
+            # dof_pos_limits = -2.0
+            # feet_regulation = -0.05
+            # hip_to_default = -0.1
+            # similar_to_default = -0.05
+
+            # CTS reward
+            tracking_lin_vel = 1.0
+            tracking_ang_vel = 0.5
+            lin_vel_z = -2.0
+            ang_vel_xy = -0.05
+            dof_acc = -2.5e-7
+            dof_power = -2e-5
+            torques = -1e-4
+            correct_base_height = -1.0
+            action_rate = -0.01
+            action_smoothness = -0.01
+            collision = -1.0
+            dof_pos_limits = -2.0
+            feet_regulation = -0.05
+            # CTS奖励训出来双脚距离非常近, 真机效果很差, 但是sim2sim能上20cm楼梯, 尝试加入hip_to_default奖励或similar_to_default奖励
+            hip_to_default = -0.05  # 在训练到y=1.5时, 双脚会明显碰撞, 为避免该问题提升hip, 效果更差, 还是保持0.05 (y最大也只到0.1了)
+            # legs_distance = -1.5  # 奖励双脚距离, 避免CTS训练出来双脚距离过近, 尝试加入后robogauge flat验证效果变差, 删除
+            # similar_to_default = -0.01
+            # feet_contact_forces = -1.0  # 尝试加入但并没有起到任何效果, 删除
+
+        turn_over_roll_threshold = math.pi / 4 # threshold on roll to use turn over rewards
+        class turn_over_scales:
+            upright = 1.0
+            # dof_acc = -2.5e-7
+            # dof_power = -2e-5
+            # action_rate = -0.001
+            # action_smoothness = -0.001
+
+    class noise(LeggedRobotCfg.noise):
+        add_noise = True
+
+class GO2CfgPPO(LeggedRobotCfgPPO):
+    class algorithm(LeggedRobotCfgPPO.algorithm):
+        entropy_coef = 0.01
+    class runner(LeggedRobotCfgPPO.runner):
+        run_name = ''
+        experiment_name = 'go2_ppo'
+        max_iterations = 100000
+        save_interval = 500
+
+class GO2CfgCTS(LeggedRobotCfgCTS):
+    class runner(LeggedRobotCfgCTS.runner):
+        num_steps_per_env = 24
+        run_name = ''
+        experiment_name = 'go2_cts'
+        max_iterations = 100000
+        save_interval = 500
+    
+    class policy(LeggedRobotCfgCTS.policy):
+        latent_dim = 32
+        norm_type = 'l2norm'
+
+class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
+    class policy(LeggedRobotCfgMoECTS.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):
+        load_balance_coef = 0.01
+
+    class runner(LeggedRobotCfgMoECTS.runner):
+        run_name = ''
+        experiment_name = 'go2_moe_cts'
+        max_iterations = 100000
+        save_interval = 500

legged_gym/envs/go2/go2_env.py

@@ -0,0 +1,53 @@
+
+from legged_gym.envs.base.legged_robot import LeggedRobot
+
+from isaacgym.torch_utils import *
+from isaacgym import gymtorch, gymapi, gymutil
+import torch
+
+class Go2Robot(LeggedRobot):
+    def _get_noise_scale_vec(self, cfg):
+        noise_vec = torch.zeros_like(self.obs_buf[0])
+        self.add_noise = self.cfg.noise.add_noise
+        noise_scales = self.cfg.noise.noise_scales
+        noise_level = self.cfg.noise.noise_level
+        noise_vec[:3] = noise_scales.ang_vel * noise_level * self.obs_scales.ang_vel
+        noise_vec[3:6] = noise_scales.gravity * noise_level
+        noise_vec[6:9] = 0. # commands
+        noise_vec[9:9+self.num_actions] = noise_scales.dof_pos * noise_level * self.obs_scales.dof_pos
+        noise_vec[9+self.num_actions:9+2*self.num_actions] = noise_scales.dof_vel * noise_level * self.obs_scales.dof_vel
+        noise_vec[9+2*self.num_actions:9+3*self.num_actions] = 0. # previous actions
+
+        return noise_vec
+    
+    def compute_observations(self):
+        """ Computes observations
+        """
+        self.obs_buf = torch.cat((self.base_ang_vel  * self.obs_scales.ang_vel,
+                                  self.projected_gravity,
+                                  self.commands[:, :3] * self.commands_scale,
+                                  (self.dof_pos - self.default_dof_pos) * self.obs_scales.dof_pos,
+                                  self.dof_vel * self.obs_scales.dof_vel,
+                                  self.actions,
+                                  ),dim=-1)
+        
+        heights = torch.clip(self.root_states[:, 2].unsqueeze(1) - 0.5 - self.measured_heights, -1, 1.0) * self.obs_scales.height_measurements
+        
+        self.privileged_obs_buf = torch.cat((  self.base_lin_vel * self.obs_scales.lin_vel,
+                                    self.base_ang_vel  * self.obs_scales.ang_vel,
+                                    self.projected_gravity,
+                                    self.commands[:, :3] * self.commands_scale,
+                                    (self.dof_pos - self.default_dof_pos) * self.obs_scales.dof_pos,
+                                    self.dof_vel * self.obs_scales.dof_vel,
+                                    self.actions,
+                                    torch.norm(self.contact_forces[:, self.feet_indices, :], dim=-1) * 1e-3,  # foot contact forces (4,)
+                                    self.torques / self.torque_limits,  # motor torques (12,)
+                                    (self.last_dof_vel - self.dof_vel) / self.dt * 1e-4,  # motor accelerations (12,)
+                                    heights,  # height measurements (187,)
+                                    ),dim=-1)
+        # print(f"foot contact: {self.privileged_obs_buf[:,48:48+4].min(), self.privileged_obs_buf[:,48:48+4].max()}")
+        # print(f"torques: {self.privileged_obs_buf[:,48+4:48+4+12].min(), self.privileged_obs_buf[:,48+4:48+4+12].max()}")
+        # print(f"acc: {self.privileged_obs_buf[:,48+4+12:48+4+12+12].min(), self.privileged_obs_buf[:,48+4+12:48+4+12+12].max()}")
+        
+        if self.add_noise:
+            self.obs_buf += (2 * torch.rand_like(self.obs_buf) - 1) * self.noise_scale_vec