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

Init

Browse Source
This commit is contained in:
wty-yy
2025-12-29 16:33:09 +08:00
parent 83952bd549
commit 0c56db7e42
96 changed files with 770085 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
legged_gym/LICENSE Normal file
View File

@@ -0,0 +1,31 @@
Copyright (c) 2021, ETH Zurich, Nikita Rudin
Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES
All rights reserved.
Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:
1. Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.
3. Neither the name of the copyright holder nor the names of its contributors
may be used to endorse or promote products derived from this software without
specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
See licenses/assets for license information for assets included in this repository.
See licenses/dependencies for license information of dependencies of this package.

4
legged_gym/__init__.py Normal file
View File

@@ -0,0 +1,4 @@
import os
LEGGED_GYM_ROOT_DIR = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
LEGGED_GYM_ENVS_DIR = os.path.join(LEGGED_GYM_ROOT_DIR, 'legged_gym', 'envs')

11
legged_gym/envs/__init__.py Normal file
View File

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

25
legged_gym/envs/base/base_config.py Normal file
View File

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

115
legged_gym/envs/base/base_task.py Normal file
View File

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

1409
legged_gym/envs/base/legged_robot.py Normal file
View File

File diff suppressed because it is too large Load Diff

361
legged_gym/envs/base/legged_robot_config.py Normal file
View File

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

266
legged_gym/envs/go2/go2_config.py Normal file
View File

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

53
legged_gym/envs/go2/go2_env.py Normal file
View File

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

70
legged_gym/scripts/play.py Normal file
View File

@@ -0,0 +1,70 @@
import sys
from legged_gym import LEGGED_GYM_ROOT_DIR
import os
import sys
from legged_gym import LEGGED_GYM_ROOT_DIR
import isaacgym
from legged_gym.envs import *
from legged_gym.utils import get_args, task_registry, Logger
from legged_gym.utils.exporter import export_policy_as_jit, export_policy_as_onnx, export_policy_as_pkl
import numpy as np
import torch
def play(args):
env_cfg, train_cfg = task_registry.get_cfgs(name=args.task)
# override some parameters for testing
env_cfg.env.num_envs = min(env_cfg.env.num_envs, 100)
# env_cfg.terrain.mesh_type = 'plane'
env_cfg.terrain.num_rows = 7
env_cfg.terrain.num_cols = 7
env_cfg.terrain.curriculum = False
env_cfg.noise.add_noise = False
env_cfg.domain_rand.randomize_friction = False
env_cfg.domain_rand.push_robots = False
env_cfg.domain_rand.randomize_base_mass = False
env_cfg.domain_rand.randomize_link_mass = False
env_cfg.domain_rand.randomize_base_com = False
env_cfg.domain_rand.randomize_pd_gains = False
env_cfg.domain_rand.randomize_motor_zero_offset = False
env_cfg.env.test = True
# prepare environment
env, _ = task_registry.make_env(name=args.task, args=args, env_cfg=env_cfg)
obs = env.get_observations()
# load policy
train_cfg.runner.resume = True
runner, train_cfg = task_registry.make_alg_runner(env=env, name=args.task, args=args, train_cfg=train_cfg)
policy = runner.get_inference_policy(device=env.device)
# export policy as a jit module (used to run it from C++)
if EXPORT_POLICY:
path = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs', train_cfg.runner.experiment_name, 'exported', 'policies')
if hasattr(runner.alg, 'actor_critic'):
model = runner.alg.actor_critic
else:
model = runner.alg.model
export_policy_as_jit(model, path)
export_policy_as_onnx(model, path)
export_policy_as_pkl(model, path)
print('Exported policy as jit script / onnx to: ', path)
for i in range(10*int(env.max_episode_length)):
actions = policy(obs.detach())
if FIX_COMMAND:
env.commands[:, 0] = 1.0
env.commands[:, 1] = 0.0
env.commands[:, 2] = 0.0
obs, _, rews, dones, infos = env.step(actions.detach())
if __name__ == '__main__':
EXPORT_POLICY = True
RECORD_FRAMES = False
MOVE_CAMERA = False
FIX_COMMAND = True
args = get_args()
play(args)

20
legged_gym/scripts/train.py Normal file
View File

@@ -0,0 +1,20 @@
import os
import numpy as np
from datetime import datetime
import sys
import isaacgym
from legged_gym.envs import *
from legged_gym.utils import get_args, task_registry
import torch
def train(args):
env, env_cfg = task_registry.make_env(name=args.task, args=args)
runner, train_cfg = task_registry.make_alg_runner(env=env, name=args.task, args=args)
env.common_step_counter = runner.current_learning_iteration * env.num_steps_per_env # resume env step counter
env.update_reward_curriculum(force_update=True) # force update reward curriculum at start
runner.learn(num_learning_iterations=train_cfg.runner.max_iterations, init_at_random_ep_len=True)
if __name__ == '__main__':
args = get_args()
train(args)

5
legged_gym/utils/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
from .helpers import class_to_dict, get_load_path, get_args, set_seed, update_class_from_dict
from .task_registry import task_registry
from .logger import Logger
from .math import *
from .terrain import Terrain

268
legged_gym/utils/exporter.py Normal file
View File

@@ -0,0 +1,268 @@
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
import copy
import os
import torch
from torch import nn
from typing import Optional
def export_policy_as_jit(policy: object, path: str, normalizer: Optional[object] = None, filename="policy.pt"):
"""Export policy into a Torch JIT file.
Args:
policy: The policy torch module.
normalizer: The empirical normalizer module. If None, Identity is used.
path: The path to the saving directory.
filename: The name of exported JIT file. Defaults to "policy.pt".
"""
policy_exporter = _TorchPolicyExporter(policy, normalizer)
policy_exporter.export(path, filename)
def export_policy_as_onnx(
policy: object, path: str, normalizer: Optional[object] = None, filename="policy.onnx", verbose=False
):
"""Export policy into a Torch ONNX file.
Args:
policy: The policy torch module.
normalizer: The empirical normalizer module. If None, Identity is used.
path: The path to the saving directory.
filename: The name of exported ONNX file. Defaults to "policy.onnx".
verbose: Whether to print the model summary. Defaults to False.
"""
if not os.path.exists(path):
os.makedirs(path, exist_ok=True)
policy_exporter = _OnnxPolicyExporter(policy, normalizer, verbose)
policy_exporter.export(path, filename)
def export_policy_as_pkl(
policy: nn.Module, path: str, filename="policy.pkl"
):
"""Export policy into a Torch pkl file.
Args:
policy: The policy torch module.
normalizer: The empirical normalizer module. If None, Identity is used.
path: The path to the saving directory.
filename: The name of exported pkl file. Defaults to "policy.pkl".
"""
if not os.path.exists(path):
os.makedirs(path, exist_ok=True)
model_dict = policy.state_dict()
torch.save(model_dict, os.path.join(path, filename))
"""
Helper Classes - Private.
"""
class _TorchPolicyExporter(torch.nn.Module):
"""Exporter of actor-critic into JIT file."""
def __init__(self, policy, normalizer=None):
super().__init__()
self.is_recurrent = policy.is_recurrent
# copy policy parameters
if hasattr(policy, "student_encoder"):
self.student_encoder = copy.deepcopy(policy.student_encoder).cpu()
self.history = torch.zeros([1, policy.history.shape[1], policy.history.shape[2]], device='cpu')
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()
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"):
self.actor = copy.deepcopy(policy.actor)
if self.is_recurrent:
self.rnn = copy.deepcopy(policy.memory_a.rnn)
elif hasattr(policy, "student"):
self.actor = copy.deepcopy(policy.student)
if self.is_recurrent:
self.rnn = copy.deepcopy(policy.memory_s.rnn)
else:
raise ValueError("Policy does not have an actor/student module.")
# set up recurrent network
if self.is_recurrent:
self.rnn.cpu()
self.register_buffer("hidden_state", torch.zeros(self.rnn.num_layers, 1, self.rnn.hidden_size))
self.register_buffer("cell_state", torch.zeros(self.rnn.num_layers, 1, self.rnn.hidden_size))
self.forward = self.forward_lstm
self.reset = self.reset_memory
# copy normalizer if exists
if normalizer:
self.normalizer = copy.deepcopy(normalizer)
else:
self.normalizer = torch.nn.Identity()
def forward_lstm(self, x):
x = self.normalizer(x)
x, (h, c) = self.rnn(x.unsqueeze(0), (self.hidden_state, self.cell_state))
self.hidden_state[:] = h
self.cell_state[:] = c
x = x.squeeze(0)
return self.actor(x)
def forward(self, x):
return self.actor(self.normalizer(x))
def forward_cts(self, x): # x is single observations
x = self.normalizer(x)
self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
latent = self.student_encoder(self.history.flatten(1))
x = torch.cat([latent, x], dim=1)
return self.actor(x)
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)
history_no_goal = self.history.reshape(1, self.history_length, -1)[:, :, self.obs_no_goal_mask].reshape(1, -1)
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)
@torch.jit.export
def reset(self):
if hasattr(self, 'history'):
self.history = torch.zeros_like(self.history)
def reset_memory(self):
self.hidden_state[:] = 0.0
self.cell_state[:] = 0.0
def export(self, path, filename):
os.makedirs(path, exist_ok=True)
path = os.path.join(path, filename)
self.to("cpu")
traced_script_module = torch.jit.script(self)
traced_script_module.save(path)
class _OnnxPolicyExporter(torch.nn.Module):
"""Exporter of actor-critic into ONNX file."""
def __init__(self, policy, normalizer=None, verbose=False):
super().__init__()
self.verbose = verbose
self.input_dim = None
self.num_actions = 12
self.normalizer = torch.nn.Identity()
# copy policy parameters
if hasattr(policy, 'student_encoder'):
self.student_encoder = copy.deepcopy(policy.student_encoder)
self.forward = self.forward_cts
self.input_dim = self.student_encoder[0].in_features
elif hasattr(policy, "student_moe_encoder"):
self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
self.history_length = policy.history.shape[1]
self.forward = self.forward_moe_cts
self.input_dim = self.history_length * policy.history.shape[2]
if hasattr(policy, "actor"):
self.actor = copy.deepcopy(policy.actor)
if hasattr(self, 'is_recurrent') and self.is_recurrent:
self.rnn = copy.deepcopy(policy.memory_a.rnn)
if self.input_dim is None:
self.input_dim = self.actor[0].in_features
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)
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)
last_obs = history[:, -obs_dim:]
latent = self.student_encoder(history)
x = torch.cat([latent, last_obs], dim=1)
return self.actor(x)
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)
last_obs = history[:, -obs_dim:]
history_3d = history.view(-1, self.history_length, obs_dim)
history_no_goal = history_3d[:, :, self.obs_no_goal_mask].reshape(x.shape[0], -1)
latent, weights = self.student_moe_encoder(history, history_no_goal)
x = torch.cat([latent, last_obs], dim=1)
return self.actor(x), weights, latent
def export(self, path, filename):
self.to("cpu")
obs = torch.zeros(1, self.input_dim)
output_names = ["actions"]
if self.forward == self.forward_moe_cts:
output_names.append("weights")
output_names.append("latent")
torch.onnx.export(
self,
obs,
os.path.join(path, filename),
export_params=True,
opset_version=11,
verbose=self.verbose,
input_names=["obs"],
output_names=output_names,
dynamic_axes={},
)

193
legged_gym/utils/helpers.py Normal file
View File

@@ -0,0 +1,193 @@
import os
import copy
import torch
import numpy as np
import random
from pathlib import Path
from isaacgym import gymapi
from isaacgym import gymutil
from legged_gym import LEGGED_GYM_ROOT_DIR, LEGGED_GYM_ENVS_DIR
def class_to_dict(obj) -> dict:
if not hasattr(obj,"__dict__"):
return obj
result = {}
for key in dir(obj):
if key.startswith("_"):
continue
element = []
val = getattr(obj, key)
if isinstance(val, list):
for item in val:
element.append(class_to_dict(item))
else:
element = class_to_dict(val)
result[key] = element
return result
def update_class_from_dict(obj, dict):
for key, val in dict.items():
attr = getattr(obj, key, None)
if isinstance(attr, type):
update_class_from_dict(attr, val)
else:
setattr(obj, key, val)
return
def set_seed(seed):
if seed == -1:
seed = np.random.randint(0, 10000)
print("Setting seed: {}".format(seed))
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def parse_sim_params(args, cfg):
# code from Isaac Gym Preview 2
# initialize sim params
sim_params = gymapi.SimParams()
# set some values from args
if args.physics_engine == gymapi.SIM_FLEX:
if args.device != "cpu":
print("WARNING: Using Flex with GPU instead of PHYSX!")
elif args.physics_engine == gymapi.SIM_PHYSX:
sim_params.physx.use_gpu = args.use_gpu
sim_params.physx.num_subscenes = args.subscenes
sim_params.use_gpu_pipeline = args.use_gpu_pipeline
# if sim options are provided in cfg, parse them and update/override above:
if "sim" in cfg:
gymutil.parse_sim_config(cfg["sim"], sim_params)
# Override num_threads if passed on the command line
if args.physics_engine == gymapi.SIM_PHYSX and args.num_threads > 0:
sim_params.physx.num_threads = args.num_threads
return sim_params
def get_load_path(root, load_run=-1, checkpoint=-1):
try:
runs = os.listdir(root)
runs = [run for run in runs if len(list((Path(root) / run).glob('model_*.pt'))) > 0]
#TODO sort by date to handle change of month
runs.sort()
if 'exported' in runs: runs.remove('exported')
last_run = os.path.join(root, runs[-1])
except:
raise ValueError("No runs in this directory: " + root)
if load_run==-1:
load_run = last_run
else:
load_run = os.path.join(root, load_run)
if checkpoint==-1:
models = [file for file in os.listdir(load_run) if 'model' in file]
models.sort(key=lambda m: '{0:0>15}'.format(m))
model = models[-1]
else:
model = "model_{}.pt".format(checkpoint)
load_path = os.path.join(load_run, model)
return load_path
def update_cfg_from_args(env_cfg, cfg_train, args):
# seed
if env_cfg is not None:
# num envs
if args.num_envs is not None:
env_cfg.env.num_envs = args.num_envs
if cfg_train is not None:
if args.seed is not None:
cfg_train.seed = args.seed
# alg runner parameters
if args.max_iterations is not None:
cfg_train.runner.max_iterations = args.max_iterations
if args.resume:
cfg_train.runner.resume = args.resume
if args.experiment_name is not None:
cfg_train.runner.experiment_name = args.experiment_name
if args.run_name is not None:
cfg_train.runner.run_name = args.run_name
if args.load_run is not None:
cfg_train.runner.load_run = args.load_run
if args.checkpoint is not None:
cfg_train.runner.checkpoint = args.checkpoint
return env_cfg, cfg_train
def get_args():
custom_parameters = [
{"name": "--task", "type": str, "default": "go2", "help": "Resume training or start testing from a checkpoint. Overrides config file if provided."},
{"name": "--resume", "action": "store_true", "default": False, "help": "Resume training from a checkpoint"},
{"name": "--experiment_name", "type": str, "help": "Name of the experiment to run or load. Overrides config file if provided."},
{"name": "--run_name", "type": str, "help": "Name of the run. Overrides config file if provided."},
{"name": "--load_run", "type": str, "help": "Name of the run to load when resume=True. If -1: will load the last run. Overrides config file if provided."},
{"name": "--checkpoint", "type": int, "help": "Saved model checkpoint number. If -1: will load the last checkpoint. Overrides config file if provided."},
{"name": "--headless", "action": "store_true", "default": False, "help": "Force display off at all times"},
{"name": "--horovod", "action": "store_true", "default": False, "help": "Use horovod for multi-gpu training"},
{"name": "--rl_device", "type": str, "default": "cuda:0", "help": 'Device used by the RL algorithm, (cpu, gpu, cuda:0, cuda:1 etc..)'},
{"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."},
]
# parse arguments
args = gymutil.parse_arguments(
description="RL Policy",
custom_parameters=custom_parameters)
# name allignment
args.sim_device_id = args.compute_device_id
args.sim_device = args.sim_device_type
if args.sim_device=='cuda':
args.sim_device += f":{args.sim_device_id}"
return args
# def export_policy_as_jit(actor_critic, path):
# if hasattr(actor_critic, 'memory_a'):
# # assumes LSTM: TODO add GRU
# exporter = PolicyExporterLSTM(actor_critic)
# exporter.export(path)
# else:
# os.makedirs(path, exist_ok=True)
# path = os.path.join(path, 'policy_1.pt')
# model = copy.deepcopy(actor_critic.actor).to('cpu')
# traced_script_module = torch.jit.script(model)
# traced_script_module.save(path)
# class PolicyExporterLSTM(torch.nn.Module):
# def __init__(self, actor_critic):
# super().__init__()
# self.actor = copy.deepcopy(actor_critic.actor)
# self.is_recurrent = actor_critic.is_recurrent
# self.memory = copy.deepcopy(actor_critic.memory_a.rnn)
# self.memory.cpu()
# self.register_buffer(f'hidden_state', torch.zeros(self.memory.num_layers, 1, self.memory.hidden_size))
# self.register_buffer(f'cell_state', torch.zeros(self.memory.num_layers, 1, self.memory.hidden_size))
# def forward(self, x):
# out, (h, c) = self.memory(x.unsqueeze(0), (self.hidden_state, self.cell_state))
# self.hidden_state[:] = h
# self.cell_state[:] = c
# return self.actor(out.squeeze(0))
# @torch.jit.export
# def reset_memory(self):
# self.hidden_state[:] = 0.
# self.cell_state[:] = 0.
# def export(self, path):
# os.makedirs(path, exist_ok=True)
# path = os.path.join(path, 'policy_lstm_1.pt')
# self.to('cpu')
# traced_script_module = torch.jit.script(self)
# traced_script_module.save(path)

47
legged_gym/utils/isaacgym_utils.py Normal file
View File

@@ -0,0 +1,47 @@
import os
import numpy as np
import random
import torch
@torch.jit.script
def copysign(a, b):
# type: (float, Tensor) -> Tensor
a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0])
return torch.abs(a) * torch.sign(b)
def get_euler_xyz(q):
qx, qy, qz, qw = 0, 1, 2, 3
# roll (x-axis rotation)
sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz])
cosr_cosp = q[:, qw] * q[:, qw] - q[:, qx] * \
q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz]
roll = torch.atan2(sinr_cosp, cosr_cosp)
# pitch (y-axis rotation)
sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx])
pitch = torch.where(
torch.abs(sinp) >= 1, copysign(np.pi / 2.0, sinp), torch.asin(sinp))
# yaw (z-axis rotation)
siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy])
cosy_cosp = q[:, qw] * q[:, qw] + q[:, qx] * \
q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz]
yaw = torch.atan2(siny_cosp, cosy_cosp)
return torch.stack((roll, pitch, yaw), dim=-1)
def sample_disjoint_intervals(env_ids, limit_bound, cfg_min, cfg_max, device):
"""
sample uniform distribution from [cfg_min, -limit_bound] U [limit_bound, cfg_max]
"""
width_neg = torch.nn.functional.relu(-limit_bound - cfg_min)
width_pos = torch.nn.functional.relu(cfg_max - limit_bound)
total_width = width_neg + width_pos + 1e-6 # 加极小值防除零
u = torch.rand(len(env_ids), device=device) * total_width
samples = torch.where(
u < width_neg,
cfg_min + u,
cfg_max - width_pos + (u - width_neg)
)
return samples

39
legged_gym/utils/logger.py Normal file
View File

@@ -0,0 +1,39 @@
import numpy as np
from collections import defaultdict
from multiprocessing import Process, Value
class Logger:
def __init__(self, dt):
self.state_log = defaultdict(list)
self.rew_log = defaultdict(list)
self.dt = dt
self.num_episodes = 0
self.plot_process = None
def log_state(self, key, value):
self.state_log[key].append(value)
def log_states(self, dict):
for key, value in dict.items():
self.log_state(key, value)
def log_rewards(self, dict, num_episodes):
for key, value in dict.items():
if 'rew' in key:
self.rew_log[key].append(value.item() * num_episodes)
self.num_episodes += num_episodes
def reset(self):
self.state_log.clear()
self.rew_log.clear()
def print_rewards(self):
print("Average rewards per second:")
for key, values in self.rew_log.items():
mean = np.sum(np.array(values)) / self.num_episodes
print(f" - {key}: {mean}")
print(f"Total number of episodes: {self.num_episodes}")
def __del__(self):
if self.plot_process is not None:
self.plot_process.kill()

26
legged_gym/utils/math.py Normal file
View File

@@ -0,0 +1,26 @@
import torch
from torch import Tensor
import numpy as np
from isaacgym.torch_utils import quat_apply, normalize
from typing import Tuple
# @ torch.jit.script
def quat_apply_yaw(quat, vec):
quat_yaw = quat.clone().view(-1, 4)
quat_yaw[:, :2] = 0.
quat_yaw = normalize(quat_yaw)
return quat_apply(quat_yaw, vec)
# @ torch.jit.script
def wrap_to_pi(angles):
angles %= 2*np.pi
angles -= 2*np.pi * (angles > np.pi)
return angles
# @ torch.jit.script
def torch_rand_sqrt_float(lower, upper, shape, device):
# type: (float, float, Tuple[int, int], str) -> Tensor
r = 2*torch.rand(*shape, device=device) - 1
r = torch.where(r<0., -torch.sqrt(-r), torch.sqrt(r))
r = (r + 1.) / 2.
return (upper - lower) * r + lower

129
legged_gym/utils/task_registry.py Normal file
View File

@@ -0,0 +1,129 @@
import os
from datetime import datetime
from typing import Tuple
import torch
import numpy as np
import sys
from rsl_rl.env import VecEnv
from rsl_rl.runners import OnPolicyRunner, OnPolicyRunnerCTS
from legged_gym import LEGGED_GYM_ROOT_DIR, LEGGED_GYM_ENVS_DIR
from .helpers import get_args, update_cfg_from_args, class_to_dict, get_load_path, set_seed, parse_sim_params
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO
class TaskRegistry():
def __init__(self):
self.task_classes = {}
self.env_cfgs = {}
self.train_cfgs = {}
def register(self, name: str, task_class: VecEnv, env_cfg: LeggedRobotCfg, train_cfg: LeggedRobotCfgPPO):
self.task_classes[name] = task_class
self.env_cfgs[name] = env_cfg
self.train_cfgs[name] = train_cfg
def get_task_class(self, name: str) -> VecEnv:
return self.task_classes[name]
def get_cfgs(self, name) -> Tuple[LeggedRobotCfg, LeggedRobotCfgPPO]:
train_cfg = self.train_cfgs[name]
env_cfg = self.env_cfgs[name]
# copy seed
env_cfg.seed = train_cfg.seed
return env_cfg, train_cfg
def make_env(self, name, args=None, env_cfg=None) -> Tuple[VecEnv, LeggedRobotCfg]:
""" Creates an environment either from a registered namme or from the provided config file.
Args:
name (string): Name of a registered env.
args (Args, optional): Isaac Gym comand line arguments. If None get_args() will be called. Defaults to None.
env_cfg (Dict, optional): Environment config file used to override the registered config. Defaults to None.
Raises:
ValueError: Error if no registered env corresponds to 'name'
Returns:
isaacgym.VecTaskPython: The created environment
Dict: the corresponding config file
"""
# if no args passed get command line arguments
if args is None:
args = get_args()
# check if there is a registered env with that name
if name in self.task_classes:
task_class = self.get_task_class(name)
else:
raise ValueError(f"Task with name: {name} was not registered")
if env_cfg is None:
# load config files
env_cfg, _ = self.get_cfgs(name)
# override cfg from args (if specified)
env_cfg, _ = update_cfg_from_args(env_cfg, None, args)
set_seed(env_cfg.seed)
# parse sim params (convert to dict first)
sim_params = {"sim": class_to_dict(env_cfg.sim)}
sim_params = parse_sim_params(args, sim_params)
env = task_class( cfg=env_cfg,
sim_params=sim_params,
physics_engine=args.physics_engine,
sim_device=args.sim_device,
headless=args.headless)
return env, env_cfg
def make_alg_runner(self, env, name=None, args=None, train_cfg=None, log_root="default") -> Tuple[OnPolicyRunner, LeggedRobotCfgPPO]:
""" Creates the training algorithm either from a registered namme or from the provided config file.
Args:
env (isaacgym.VecTaskPython): The environment to train (TODO: remove from within the algorithm)
name (string, optional): Name of a registered env. If None, the config file will be used instead. Defaults to None.
args (Args, optional): Isaac Gym comand line arguments. If None get_args() will be called. Defaults to None.
train_cfg (Dict, optional): Training config file. If None 'name' will be used to get the config file. Defaults to None.
log_root (str, optional): Logging directory for Tensorboard. Set to 'None' to avoid logging (at test time for example).
Logs will be saved in <log_root>/<date_time>_<run_name>. Defaults to "default"=<path_to_LEGGED_GYM>/logs/<experiment_name>.
Raises:
ValueError: Error if neither 'name' or 'train_cfg' are provided
Warning: If both 'name' or 'train_cfg' are provided 'name' is ignored
Returns:
PPO: The created algorithm
Dict: the corresponding config file
"""
# if no args passed get command line arguments
if args is None:
args = get_args()
# if config files are passed use them, otherwise load from the name
if train_cfg is None:
if name is None:
raise ValueError("Either 'name' or 'train_cfg' must be not None")
# load config files
_, train_cfg = self.get_cfgs(name)
else:
if name is not None:
print(f"'train_cfg' provided -> Ignoring 'name={name}'")
# override cfg from args (if specified)
_, train_cfg = update_cfg_from_args(None, train_cfg, args)
if log_root=="default":
log_root = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs', train_cfg.runner.experiment_name)
log_dir = os.path.join(log_root, datetime.now().strftime('%b%d_%H-%M-%S') + '_' + train_cfg.runner.run_name)
elif log_root is None:
log_dir = None
else:
log_dir = os.path.join(log_root, datetime.now().strftime('%b%d_%H-%M-%S') + '_' + train_cfg.runner.run_name)
train_cfg_dict = class_to_dict(train_cfg)
runner = eval(train_cfg.runner_class_name)(env, train_cfg_dict, log_dir, device=args.rl_device)
#save resume path before creating a new log_dir
resume = train_cfg.runner.resume
if resume:
# load previously trained model
resume_path = get_load_path(log_root, load_run=train_cfg.runner.load_run, checkpoint=train_cfg.runner.checkpoint)
print(f"Loading model from: {resume_path}")
runner.load(resume_path)
return runner, train_cfg
# make global task registry
task_registry = TaskRegistry()

188
legged_gym/utils/terrain.py Normal file
View File

@@ -0,0 +1,188 @@
from collections import defaultdict
import numpy as np
from numpy.random import choice
from scipy import interpolate
from isaacgym import terrain_utils
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg
class Terrain:
def __init__(self, cfg: LeggedRobotCfg.terrain, num_robots) -> None:
self.cfg = cfg
self.num_robots = num_robots
self.type = cfg.mesh_type
if self.type in ["none", 'plane']:
return
self.env_length = cfg.terrain_length
self.env_width = cfg.terrain_width
self.proportions = [np.sum(cfg.terrain_proportions[:i+1]) for i in range(len(cfg.terrain_proportions))]
self.cfg.num_sub_terrains = cfg.num_rows * cfg.num_cols
self.env_origins = np.zeros((cfg.num_rows, cfg.num_cols, 3))
self.width_per_env_pixels = int(self.env_width / cfg.horizontal_scale)
self.length_per_env_pixels = int(self.env_length / cfg.horizontal_scale)
self.spacing = cfg.terrain_spacing
self.spacing_pixels = int(self.spacing / cfg.horizontal_scale)
self.border = int(cfg.border_size/self.cfg.horizontal_scale)
self.tot_cols = int(cfg.num_cols * self.width_per_env_pixels + max(0, cfg.num_cols-1) * self.spacing_pixels) + 2 * self.border
self.tot_rows = int(cfg.num_rows * self.length_per_env_pixels + max(0, cfg.num_rows-1) * self.spacing_pixels) + 2 * self.border
self.name2cols = defaultdict(set) # terrain type to column index
self.cols2id = [] # column index to terrain id
self.height_field_raw = np.zeros((self.tot_rows , self.tot_cols), dtype=np.int16)
if cfg.curriculum:
self.curiculum()
elif cfg.selected:
self.selected_terrain()
else:
self.randomized_terrain()
self.heightsamples = self.height_field_raw
if self.type=="trimesh":
self.vertices, self.triangles = terrain_utils.convert_heightfield_to_trimesh( self.height_field_raw,
self.cfg.horizontal_scale,
self.cfg.vertical_scale,
self.cfg.slope_treshold)
def randomized_terrain(self):
for k in range(self.cfg.num_sub_terrains):
# Env coordinates in the world
(i, j) = np.unravel_index(k, (self.cfg.num_rows, self.cfg.num_cols))
choice = np.random.uniform(0, 1)
difficulty = np.random.choice([0.5, 0.75, 0.9])
terrain = self.make_terrain(choice, difficulty)
self.add_terrain_to_map(terrain, i, j)
def curiculum(self):
for j in range(self.cfg.num_cols):
for i in range(self.cfg.num_rows):
difficulty = i / self.cfg.num_rows
choice = j / self.cfg.num_cols + 0.001
terrain = self.make_terrain(choice, difficulty)
self.add_terrain_to_map(terrain, i, j)
self.name2cols[terrain.terrain_name].add(j)
self.cols2id.append(terrain.terrain_id)
def selected_terrain(self):
terrain_type = self.cfg.terrain_kwargs.pop('type')
for k in range(self.cfg.num_sub_terrains):
# Env coordinates in the world
(i, j) = np.unravel_index(k, (self.cfg.num_rows, self.cfg.num_cols))
terrain = terrain_utils.SubTerrain("terrain",
width=self.width_per_env_pixels,
length=self.width_per_env_pixels,
vertical_scale=self.vertical_scale,
horizontal_scale=self.horizontal_scale)
eval(terrain_type)(terrain, **self.cfg.terrain_kwargs.terrain_kwargs)
self.add_terrain_to_map(terrain, i, j)
def make_terrain(self, choice, difficulty):
terrain = terrain_utils.SubTerrain("terrain",
width=self.width_per_env_pixels,
length=self.width_per_env_pixels,
vertical_scale=self.cfg.vertical_scale,
horizontal_scale=self.cfg.horizontal_scale)
slope = 0.1 + difficulty * 0.52
step_height = 0.05 + 0.23 * difficulty
discrete_obstacles_height = 0.05 + difficulty * 0.25
stepping_stones_size = 1.5 * (1.05 - difficulty)
stone_distance = 0.05 if difficulty==0 else 0.1
gap_size = 1. * difficulty
# pit_depth = 1. * difficulty # 删除坑洞
amplitude = 0.1 + 0.2 * difficulty
if choice < self.proportions[0]:
terrain.terrain_name = "wave"
terrain.terrain_id = 0
terrain_utils.wave_terrain(terrain, num_waves=5, amplitude=amplitude)
terrain_utils.random_uniform_terrain(terrain, min_height=-0.05, max_height=0.05, step=0.005, downsampled_scale=0.2)
elif choice < self.proportions[1]: # 平滑坡
terrain.terrain_name = "slope"
terrain.terrain_id = 1
if choice < (self.proportions[0] + self.proportions[1])/ 2: # 一半正坡, 一半负坡
slope *= -1
terrain_utils.pyramid_sloped_terrain(terrain, slope=slope, platform_size=3.)
elif choice < self.proportions[2]: # 粗糙坡
terrain.terrain_name = "rough_slope"
terrain.terrain_id = 2
terrain_utils.pyramid_sloped_terrain(terrain, slope=slope, platform_size=3.)
terrain_utils.random_uniform_terrain(terrain, min_height=-0.05, max_height=0.05, step=0.005, downsampled_scale=0.2)
elif choice < self.proportions[4]: # 下楼梯
terrain.terrain_name = "stairs_down"
terrain.terrain_id = 4
if choice<self.proportions[3]: # 上楼梯
terrain.terrain_name = "stairs_up"
terrain.terrain_id = 3
step_height *= -1
terrain_utils.pyramid_stairs_terrain(terrain, step_width=0.31, step_height=step_height, platform_size=3.)
elif choice < self.proportions[5]: # 障碍物
terrain.terrain_name = "obstacles"
terrain.terrain_id = 5
num_rectangles = 20
rectangle_min_size = 1.
rectangle_max_size = 2.
terrain_utils.discrete_obstacles_terrain(terrain, discrete_obstacles_height, rectangle_min_size, rectangle_max_size, num_rectangles, platform_size=3.)
elif choice < self.proportions[6]: # 梅花桩
terrain.terrain_name = "stepping_stones"
terrain.terrain_id = 6
terrain_utils.stepping_stones_terrain(terrain, stone_size=stepping_stones_size, stone_distance=stone_distance, max_height=0., platform_size=4.)
elif choice < self.proportions[7]: # 间隙
terrain.terrain_name = "gap"
terrain.terrain_id = 7
gap_terrain(terrain, gap_size=gap_size, platform_size=3.)
else: # 平地
terrain.terrain_name = "flat"
terrain.terrain_id = 8
pit_terrain(terrain, depth=0.0, platform_size=4.)
return terrain
def add_terrain_to_map(self, terrain, row, col):
i = row
j = col
# map coordinate system
start_x = self.border + i * (self.length_per_env_pixels + self.spacing_pixels)
end_x = start_x + self.length_per_env_pixels
start_y = self.border + j * (self.width_per_env_pixels + self.spacing_pixels)
end_y = start_y + self.width_per_env_pixels
self.height_field_raw[start_x: end_x, start_y:end_y] = terrain.height_field_raw
env_origin_x = (i + 0.5) * self.env_length + i * self.spacing
env_origin_y = (j + 0.5) * self.env_width + j * self.spacing
x1 = int((self.env_length/2. - 1) / terrain.horizontal_scale)
x2 = int((self.env_length/2. + 1) / terrain.horizontal_scale)
y1 = int((self.env_width/2. - 1) / terrain.horizontal_scale)
y2 = int((self.env_width/2. + 1) / terrain.horizontal_scale)
env_origin_z = np.max(terrain.height_field_raw[x1:x2, y1:y2])*terrain.vertical_scale
self.env_origins[i, j] = [env_origin_x, env_origin_y, env_origin_z]
def gap_terrain(terrain, gap_size, platform_size=1.):
gap_size = int(gap_size / terrain.horizontal_scale)
platform_size = int(platform_size / terrain.horizontal_scale)
center_x = terrain.length // 2
center_y = terrain.width // 2
x1 = (terrain.length - platform_size) // 2
x2 = x1 + gap_size
y1 = (terrain.width - platform_size) // 2
y2 = y1 + gap_size
terrain.height_field_raw[center_x-x2 : center_x + x2, center_y-y2 : center_y + y2] = -1000
terrain.height_field_raw[center_x-x1 : center_x + x1, center_y-y1 : center_y + y1] = 0
def pit_terrain(terrain, depth, platform_size=1.):
depth = int(depth / terrain.vertical_scale)
platform_size = int(platform_size / terrain.horizontal_scale / 2)
x1 = terrain.length // 2 - platform_size
x2 = terrain.length // 2 + platform_size
y1 = terrain.width // 2 - platform_size
y2 = terrain.width // 2 + platform_size
terrain.height_field_raw[x1:x2, y1:y2] = -depth