from itertools import product from legged_gym import LEGGED_GYM_ROOT_DIR, envs import time from warnings import WarningMessage import numpy as np import os from isaacgym.torch_utils import * from isaacgym import gymtorch, gymapi, gymutil import torch from torch import Tensor from typing import Tuple, Dict from legged_gym import LEGGED_GYM_ROOT_DIR from legged_gym.envs.base.base_task import BaseTask from legged_gym.utils.math import wrap_to_pi, quat_apply_yaw from legged_gym.utils.isaacgym_utils import get_euler_xyz as get_euler_xyz_in_tensor from legged_gym.utils.isaacgym_utils import sample_disjoint_intervals from legged_gym.utils.helpers import class_to_dict from .legged_robot_config import LeggedRobotCfg from legged_gym.utils.terrain import Terrain class LeggedRobot(BaseTask): def __init__(self, cfg: LeggedRobotCfg, sim_params, physics_engine, sim_device, headless): """ Parses the provided config file, calls create_sim() (which creates, simulation and environments), initilizes pytorch buffers used during training Args: cfg (Dict): Environment config file sim_params (gymapi.SimParams): simulation parameters physics_engine (gymapi.SimType): gymapi.SIM_PHYSX (must be PhysX) device_type (string): 'cuda' or 'cpu' device_id (int): 0, 1, ... headless (bool): Run without rendering if True """ self.cfg = cfg self.sim_params = sim_params self.height_samples = None self.debug_viz = False self.init_done = False self._parse_cfg(self.cfg) super().__init__(self.cfg, sim_params, physics_engine, sim_device, headless) if not self.headless: self.set_camera(self.cfg.viewer.pos, self.cfg.viewer.lookat) self._init_buffers() self._prepare_reward_function() self.init_done = True self.reward_curriculum_scales = {} self.reward_curriculum_configs = [] if hasattr(self.cfg.rewards, "curriculum_rewards") and self.cfg.rewards.curriculum_rewards is not None: self.reward_curriculum_configs = self.cfg.rewards.curriculum_rewards for config in self.reward_curriculum_configs: self.reward_curriculum_scales[config['reward_name']] = config['start_value'] self.num_steps_per_env = 24 # PPO default num_steps_per_env def step(self, actions): """ Apply actions, simulate, call self.post_physics_step() Args: actions (torch.Tensor): Tensor of shape (num_envs, num_actions_per_env) """ clip_actions = self.cfg.normalization.clip_actions self.actions = torch.clip(actions, -clip_actions, clip_actions).to(self.device) # step physics and render each frame self.render() if self.cfg.domain_rand.randomize_action_delay: actions_start_decimation = torch.randint(0, self.cfg.control.decimation+1, (self.num_envs, 1), device=self.device) for i in range(self.cfg.control.decimation): if self.cfg.domain_rand.randomize_action_delay: use_actions = (i >= actions_start_decimation).float() input_actions = (1 - use_actions) * self.last_actions + use_actions * self.actions else: input_actions = self.actions self.torques = self._compute_torques(input_actions).view(self.torques.shape) if self.cfg.domain_rand.randomize_motor_strength: self.torques *= self.motor_strengths self.gym.set_dof_actuation_force_tensor(self.sim, gymtorch.unwrap_tensor(self.torques)) self.gym.simulate(self.sim) if self.cfg.env.test: elapsed_time = self.gym.get_elapsed_time(self.sim) sim_time = self.gym.get_sim_time(self.sim) if sim_time-elapsed_time>0: time.sleep(sim_time-elapsed_time) if self.device == 'cpu': self.gym.fetch_results(self.sim, True) self.gym.refresh_dof_state_tensor(self.sim) self.post_physics_step() # return clipped obs, clipped states (None), rewards, dones and infos clip_obs = self.cfg.normalization.clip_observations self.obs_buf = torch.clip(self.obs_buf, -clip_obs, clip_obs) if self.privileged_obs_buf is not None: self.privileged_obs_buf = torch.clip(self.privileged_obs_buf, -clip_obs, clip_obs) return self.obs_buf, self.privileged_obs_buf, self.rew_buf, self.reset_buf, self.extras def post_physics_step(self): """ check terminations, compute observations and rewards calls self._post_physics_step_callback() for common computations calls self._draw_debug_vis() if needed """ self.gym.refresh_actor_root_state_tensor(self.sim) self.gym.refresh_net_contact_force_tensor(self.sim) self.gym.refresh_rigid_body_state_tensor(self.sim) self.episode_length_buf += 1 self.common_step_counter += 1 self.commands_resampling_step -= 1 if self.cfg.init_state.turn_over: self.turn_over_timer = (self.turn_over_timer - self.dt).clip(min=0.0) self.update_reward_curriculum() # prepare quantities self.base_pos[:] = self.root_states[:, 0:3] self.base_quat[:] = self.root_states[:, 3:7] self.rpy[:] = get_euler_xyz_in_tensor(self.base_quat[:]) self.base_lin_vel[:] = quat_rotate_inverse(self.base_quat, self.root_states[:, 7:10]) self.base_ang_vel[:] = quat_rotate_inverse(self.base_quat, self.root_states[:, 10:13]) self.projected_gravity[:] = quat_rotate_inverse(self.base_quat, self.gravity_vec) self.max_move_distance = self.max_move_distance.maximum(torch.norm(self.root_states[:, :2] - self.env_origins[:, :2], dim=1)) self._post_physics_step_callback() # compute observations, rewards, resets, ... self.check_termination() self.compute_reward() env_ids = self.reset_buf.nonzero(as_tuple=False).flatten() self.reset_idx(env_ids) if self.cfg.domain_rand.push_robots: self._push_robots() self.compute_observations() # in some cases a simulation step might be required to refresh some obs (for example body positions) self.last_actions[:] = self.actions[:] self.last_dof_vel[:] = self.dof_vel[:] self.last_root_vel[:] = self.root_states[:, 7:13] def update_reward_curriculum(self, force_update: bool = False): # update reward curriculum if self.reward_curriculum_configs: if self.common_step_counter % self.num_steps_per_env == 0 or force_update: for config in self.reward_curriculum_configs: current_scale = self.get_current_scale(config) reward_name = config['reward_name'] self.reward_curriculum_scales[reward_name] = current_scale def get_current_scale(self, config): """ config: Dict {'start_iter': 0, 'end_iter': 1500, 'start_value': 1.0, 'end_value': 0.0} """ current_iter = self.common_step_counter // self.num_steps_per_env cfg_start_iter = config['start_iter'] cfg_end_iter = config['end_iter'] cfg_start_val = config['start_value'] cfg_end_val = config['end_value'] percentage = (current_iter - cfg_start_iter) / (cfg_end_iter - cfg_start_iter) percentage = max(min(percentage, 1.0), 0.0) current_scale = (1.0 - percentage) * cfg_start_val + percentage * cfg_end_val return current_scale def check_termination(self): """ Check if environments need to be reset """ self.reset_buf = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) if not self.cfg.init_state.turn_over: self.reset_buf = torch.any(torch.norm(self.contact_forces[:, self.termination_contact_indices, :], dim=-1) > 1., dim=1) # self.reset_buf |= torch.logical_or(torch.abs(self.rpy[:,1])>1.0, torch.abs(self.rpy[:,0])>0.8) self.time_out_buf = self.episode_length_buf > self.max_episode_length # no terminal reward for time-outs self.reset_buf |= self.time_out_buf def reset_idx(self, env_ids): """ Reset some environments. Calls self._reset_dofs(env_ids), self._reset_root_states(env_ids), and self._resample_commands(env_ids) [Optional] calls self._update_terrain_curriculum(env_ids), self.update_command_curriculum(env_ids) and Logs episode info Resets some buffers Args: env_ids (list[int]): List of environment ids which must be reset """ if len(env_ids) == 0: return ### Domain randomizations ### # randomization of the motor strength if self.cfg.domain_rand.randomize_motor_strength: rng = self.cfg.domain_rand.motor_strength_range self.motor_strengths[env_ids] = torch_rand_float( rng[0], rng[1], (len(env_ids), self.num_actions), device=self.device ) # randomization of the motor zero calibration for real machine if self.cfg.domain_rand.randomize_motor_zero_offset: self.motor_zero_offsets[env_ids] = torch_rand_float(self.cfg.domain_rand.motor_zero_offset_range[0], self.cfg.domain_rand.motor_zero_offset_range[1], (len(env_ids), self.num_actions), device=self.device) # randomization of the motor pd gains if self.cfg.domain_rand.randomize_pd_gains: self.p_gains_multiplier[env_ids] = torch_rand_float(self.cfg.domain_rand.stiffness_multiplier_range[0], self.cfg.domain_rand.stiffness_multiplier_range[1], (len(env_ids), self.num_actions), device=self.device) self.d_gains_multiplier[env_ids] = torch_rand_float(self.cfg.domain_rand.damping_multiplier_range[0], self.cfg.domain_rand.damping_multiplier_range[1], (len(env_ids), self.num_actions), device=self.device) # update terrain curriculum before reset root states if self.cfg.terrain.curriculum: self._update_terrain_curriculum(env_ids) # reset robot states self._reset_dofs(env_ids) self._reset_root_states(env_ids) # reset buffers self.actions[env_ids] = 0. self.last_actions[env_ids] = 0. self.last_dof_vel[env_ids] = 0. self.feet_air_time[env_ids] = 0. self.episode_length_buf[env_ids] = 0 self.reset_buf[env_ids] = 1 self.commands_resampling_step[env_ids] = self.cfg.commands.resampling_time / self.dt self.commands_xy_accumulation[env_ids] = 0.0 if self.cfg.commands.curriculum: self.update_command_curriculum(env_ids) self._resample_commands(env_ids) # fill extras self.extras["episode"] = {} if self.cfg.terrain.mesh_type in ["heightfield", "trimesh"]: self.extras["episode"]['terrain_level_all'] = torch.mean(self.terrain_levels.float()) for name, cols in self.terrain.name2cols.items(): if isinstance(cols, set): cols = self.terrain.name2cols[name] = torch.tensor(list(cols), device=self.device) self.extras["episode"]['terrain_level_' + name] = torch.mean(self.terrain_levels[torch.isin(self.terrain_types, cols)].float()) else: self.extras["episode"]['terrain_level_all'] = 0.0 for key in self.episode_sums.keys(): self.extras["episode"]['rew_' + key] = torch.mean(self.episode_sums[key][env_ids]) / self.max_episode_length_s self.episode_sums[key][env_ids] = 0. if self.cfg.commands.curriculum: self.extras["episode"]["max_command_x"] = self.command_ranges["lin_vel_x"][1] # send timeout info to the algorithm if self.cfg.env.send_timeouts: self.extras["time_outs"] = self.time_out_buf def compute_reward(self): """ Compute rewards Calls each reward function which had a non-zero scale (processed in self._prepare_reward_function()) adds each terms to the episode sums and to the total reward """ self.rew_buf[:] = 0. for i in range(len(self.reward_functions)): name = self.reward_names[i] raw_rew = self.reward_functions[i]() rew = raw_rew * self.reward_scales.get(name, 0.0) if self.cfg.init_state.turn_over: turn_over_rew = raw_rew * self.reward_turn_over_scales.get(name, 0.0) if name in self.reward_curriculum_scales: rew *= self.reward_curriculum_scales[name] if self.cfg.init_state.turn_over: turn_over_rew *= self.reward_curriculum_scales[name] if self.cfg.init_state.turn_over: need_turn_over = self.rpy[:, 0].abs() > self.cfg.rewards.turn_over_roll_threshold rew = torch.where(need_turn_over, turn_over_rew, rew) self.rew_buf += rew self.episode_sums[name] += rew if self.cfg.rewards.only_positive_rewards: self.rew_buf[:] = torch.clip(self.rew_buf[:], min=0.) # add termination reward after clipping if "termination" in self.reward_scales: rew = self._reward_termination() * self.reward_scales["termination"] self.rew_buf += rew self.episode_sums["termination"] += rew def compute_observations(self): """ Computes observations """ self.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 ),dim=-1) # add perceptive inputs if not blind # add noise if needed if self.add_noise: self.obs_buf += (2 * torch.rand_like(self.obs_buf) - 1) * self.noise_scale_vec def create_sim(self): """ Creates simulation, terrain and evironments """ self.up_axis_idx = 2 # 2 for z, 1 for y -> adapt gravity accordingly self.sim = self.gym.create_sim(self.sim_device_id, self.graphics_device_id, self.physics_engine, self.sim_params) mesh_type = self.cfg.terrain.mesh_type if mesh_type in ['heightfield', 'trimesh']: self.terrain = Terrain(self.cfg.terrain, self.num_envs) if mesh_type=='plane': self._create_ground_plane() elif mesh_type=='heightfield': self._create_heightfield() elif mesh_type=='trimesh': self._create_trimesh() elif mesh_type is not None: raise ValueError("Terrain mesh type not recognised. Allowed types are [None, plane, heightfield, trimesh]") self._create_envs() def set_camera(self, position, lookat): """ Set camera position and direction """ cam_pos = gymapi.Vec3(position[0], position[1], position[2]) cam_target = gymapi.Vec3(lookat[0], lookat[1], lookat[2]) self.gym.viewer_camera_look_at(self.viewer, None, cam_pos, cam_target) #------------- Callbacks -------------- def _process_rigid_shape_props(self, props, env_id): """ Callback allowing to store/change/randomize the rigid shape properties of each environment. Called During environment creation. Base behavior: randomizes the friction of each environment Args: props (List[gymapi.RigidShapeProperties]): Properties of each shape of the asset env_id (int): Environment id Returns: [List[gymapi.RigidShapeProperties]]: Modified rigid shape properties """ if self.cfg.domain_rand.randomize_friction: if env_id==0: # prepare friction randomization friction_range = self.cfg.domain_rand.friction_range num_buckets = 64 bucket_ids = torch.randint(0, num_buckets, (self.num_envs, 1)) friction_buckets = torch_rand_float(friction_range[0], friction_range[1], (num_buckets,1), device='cpu') self.friction_coeffs = friction_buckets[bucket_ids] for s in range(len(props)): props[s].friction = self.friction_coeffs[env_id] if self.cfg.domain_rand.randomize_restitution: rand_restitution = np.random.uniform(self.cfg.domain_rand.restitution_range[0], self.cfg.domain_rand.restitution_range[1]) for s in range(len(props)): props[s].restitution = rand_restitution return props def _process_dof_props(self, props, env_id): """ Callback allowing to store/change/randomize the DOF properties of each environment. Called During environment creation. Base behavior: stores position, velocity and torques limits defined in the URDF Args: props (numpy.array): Properties of each DOF of the asset env_id (int): Environment id Returns: [numpy.array]: Modified DOF properties """ if env_id==0: self.dof_pos_limits = torch.zeros(self.num_dof, 2, dtype=torch.float, device=self.device, requires_grad=False) self.dof_vel_limits = torch.zeros(self.num_dof, dtype=torch.float, device=self.device, requires_grad=False) self.torque_limits = torch.zeros(self.num_dof, dtype=torch.float, device=self.device, requires_grad=False) for i in range(len(props)): self.dof_pos_limits[i, 0] = props["lower"][i].item() self.dof_pos_limits[i, 1] = props["upper"][i].item() self.dof_vel_limits[i] = props["velocity"][i].item() self.torque_limits[i] = props["effort"][i].item() # soft limits m = (self.dof_pos_limits[i, 0] + self.dof_pos_limits[i, 1]) / 2 r = self.dof_pos_limits[i, 1] - self.dof_pos_limits[i, 0] self.dof_pos_limits[i, 0] = m - 0.5 * r * self.cfg.rewards.soft_dof_pos_limit self.dof_pos_limits[i, 1] = m + 0.5 * r * self.cfg.rewards.soft_dof_pos_limit return props def _process_rigid_body_props(self, props, env_id): # if env_id==0: # sum = 0 # for i, p in enumerate(props): # sum += p.mass # print(f"Mass of body {i}: {p.mass} (before randomization)") # print(f"Total mass {sum} (before randomization)") # randomize base mass if self.cfg.domain_rand.randomize_base_mass: rng = self.cfg.domain_rand.added_mass_range props[0].mass += np.random.uniform(rng[0], rng[1]) # randomize link masses if self.cfg.domain_rand.randomize_link_mass: self.multiplied_link_masses_ratio = torch_rand_float(self.cfg.domain_rand.multiplied_link_mass_range[0], self.cfg.domain_rand.multiplied_link_mass_range[1], (1, self.num_bodies-1), device=self.device) for i in range(1, len(props)): props[i].mass *= self.multiplied_link_masses_ratio[0,i-1] # randomize base com if self.cfg.domain_rand.randomize_base_com: self.added_base_com = torch_rand_float(self.cfg.domain_rand.added_base_com_range[0], self.cfg.domain_rand.added_base_com_range[1], (1, 3), device=self.device) props[0].com += gymapi.Vec3(self.added_base_com[0, 0], self.added_base_com[0, 1], self.added_base_com[0, 2]) return props def _post_physics_step_callback(self): """ Callback called before computing terminations, rewards, and observations Default behaviour: Compute ang vel command based on target and heading, compute measured terrain heights and randomly push robots """ # env_ids = (self.episode_length_buf % int(self.cfg.commands.resampling_time / self.dt)==0).nonzero(as_tuple=False).flatten() resampling_env_ids = ((self.commands_resampling_step <= 0.0) * (self.episode_length_buf < self.max_episode_length - 1)).nonzero(as_tuple=False).flatten() self._resample_commands(resampling_env_ids) if self.cfg.commands.heading_command: mask = (self.stop_heading == 0.0) forward = quat_apply(self.base_quat[mask], self.forward_vec[mask]) heading = torch.atan2(forward[:, 1], forward[:, 0]) self.commands[mask, 2] = torch.clip( 0.5*wrap_to_pi(self.commands[mask, 3] - heading), self.env_command_ranges["ang_vel_yaw"][:, 0], self.env_command_ranges["ang_vel_yaw"][:, 1] ) if self.cfg.terrain.measure_heights: self.measured_heights = self._get_heights() def _resample_commands(self, env_ids): """ Randommly select commands of some environments Args: env_ids (List[int]): Environments ids for which new commands are needed """ self.stop_heading[env_ids] = False # update command curriculum with train steps if len(self.cfg.commands.command_range_curriculum): current_iter = self.common_step_counter // self.num_steps_per_env for i in range(len(self.cfg.commands.command_range_curriculum)-1, -1, -1): # iterate backwards to be able to pop entries cfg = self.cfg.commands.command_range_curriculum[i] if current_iter >= cfg["iter"]: self.command_ranges["lin_vel_x"] = cfg["lin_vel_x"] self.command_ranges["lin_vel_y"] = cfg["lin_vel_y"] self.command_ranges["ang_vel_yaw"] = cfg["ang_vel_yaw"] self.command_ranges["heading"] = cfg["heading"] self.max_lin_vel = max(abs(self.command_ranges["lin_vel_x"][0]), abs(self.command_ranges["lin_vel_x"][1]), abs(self.command_ranges["lin_vel_y"][0]), abs(self.command_ranges["lin_vel_y"][1])) self.cfg.commands.command_range_curriculum.pop(i) self._update_env_command_ranges() print(f"Command range updated at iter {current_iter}: {self.command_ranges}") remaining_dist = torch.clip(0.625 * self.cfg.terrain.terrain_length - torch.norm(self.commands_xy_accumulation[env_ids], dim=1) * self.cfg.commands.resampling_time, 0.0) if self.cfg.commands.dynamic_resample_commands: # arrive at boundary 0.625 times the width of the remaining distance if ((self.max_episode_length - self.episode_length_buf[env_ids]) == 0).any(): raise ValueError("Some envs have zero remaining episode length during command resampling") vel_low_bound = torch.clip(remaining_dist / ((self.max_episode_length - self.episode_length_buf[env_ids] + 1e-9) * self.dt), 0.0) self.commands[env_ids, 0] = sample_disjoint_intervals( env_ids, vel_low_bound, self.env_command_ranges["lin_vel_x"][env_ids, 0], self.env_command_ranges["lin_vel_x"][env_ids, 1], self.device ) self.commands[env_ids, 1] = sample_disjoint_intervals( env_ids, vel_low_bound, self.env_command_ranges["lin_vel_y"][env_ids, 0], self.env_command_ranges["lin_vel_y"][env_ids, 1], self.device ) if self.cfg.commands.heading_command: r = torch.rand(len(env_ids), device=self.device) lower = self.env_command_ranges["heading"][env_ids, 0] upper = self.env_command_ranges["heading"][env_ids, 1] self.commands[env_ids, 3] = (upper - lower) * r + lower else: r = torch.rand(len(env_ids), device=self.device) lower = self.env_command_ranges["ang_vel_yaw"][env_ids, 0] upper = self.env_command_ranges["ang_vel_yaw"][env_ids, 1] self.commands[env_ids, 2] = (upper - lower) * r + lower self.commands_resampling_step[env_ids] = self.cfg.commands.resampling_time / self.dt else: self.commands[env_ids, 0] = torch_rand_float(self.command_ranges["lin_vel_x"][0], self.command_ranges["lin_vel_x"][1], (len(env_ids), 1), device=self.device).squeeze(1) self.commands[env_ids, 1] = torch_rand_float(self.command_ranges["lin_vel_y"][0], self.command_ranges["lin_vel_y"][1], (len(env_ids), 1), device=self.device).squeeze(1) if self.cfg.commands.heading_command: self.commands[env_ids, 3] = torch_rand_float(self.command_ranges["heading"][0], self.command_ranges["heading"][1], (len(env_ids), 1), device=self.device).squeeze(1) else: self.commands[env_ids, 2] = torch_rand_float(self.command_ranges["ang_vel_yaw"][0], self.command_ranges["ang_vel_yaw"][1], (len(env_ids), 1), device=self.device).squeeze(1) # set small commands to zero self.commands[env_ids, :2] *= (torch.norm(self.commands[env_ids, :2], dim=1) > 0.2).unsqueeze(1) # set small commands to zero # self.commands[env_ids, :2] *= (torch.norm(self.commands[env_ids, :2], dim=1) > 0.2).unsqueeze(1) rand_prob = torch.rand(len(env_ids), device=self.device) min_prob, max_prob = 0.0, 0.0 # set limitation lin vel if self.limit_vel_prob > 0.0: max_prob += self.limit_vel_prob lim_mask = (rand_prob >= min_prob) * (rand_prob < max_prob) lim_env_ids = env_ids[lim_mask] if len(lim_env_ids) > 0: change_lim_env_ids = lim_env_ids if self.cfg.commands.limit_vel_invert_when_continuous: was_limited = self.last_is_limit_vel[lim_env_ids] invert_env_ids = lim_env_ids[was_limited] self.commands[invert_env_ids, 0] *= -1.0 self.commands[invert_env_ids, 1] *= -1.0 self.commands[invert_env_ids, 2] *= -1.0 change_lim_env_ids = lim_env_ids[~was_limited] vel_idx = torch.randint(0, self.limit_vel_comb.shape[0], (len(change_lim_env_ids),), device=self.device) lin_vel_x_lim = torch.where( self.limit_vel_comb[vel_idx, 0] == -1, self.env_command_ranges["lin_vel_x"][change_lim_env_ids, 0], self.env_command_ranges["lin_vel_x"][change_lim_env_ids, 1], ) lin_vel_x_lim[self.limit_vel_comb[vel_idx, 0] == 0] = 0.0 lin_vel_y_lim = torch.where( self.limit_vel_comb[vel_idx, 1] == -1, self.env_command_ranges["lin_vel_y"][change_lim_env_ids, 0], self.env_command_ranges["lin_vel_y"][change_lim_env_ids, 1] ) lin_vel_y_lim[self.limit_vel_comb[vel_idx, 1] == 0] = 0.0 ang_vel_z_lim = torch.where( self.limit_vel_comb[vel_idx, 2] == -1, self.env_command_ranges["ang_vel_yaw"][change_lim_env_ids, 0], self.env_command_ranges["ang_vel_yaw"][change_lim_env_ids, 1] ) ang_vel_z_lim[self.limit_vel_comb[vel_idx, 2] == 0] = 0.0 self.commands[change_lim_env_ids, 0] = lin_vel_x_lim self.commands[change_lim_env_ids, 1] = lin_vel_y_lim self.commands[change_lim_env_ids, 2] = ang_vel_z_lim if self.cfg.commands.heading_command and self.cfg.commands.stop_heading_at_limit: self.stop_heading[lim_env_ids] = True # stop heading to current heading self.last_is_limit_vel[env_ids] = False self.last_is_limit_vel[lim_env_ids] = True else: self.last_is_limit_vel[env_ids] = False min_prob += self.limit_vel_prob # set all commands to zero with some probability if self.cfg.commands.zero_command_curriculum is not None: self.zero_command_proba = self.get_current_scale(self.cfg.commands.zero_command_curriculum) if self.zero_command_proba > 0.0: max_prob += self.zero_command_proba next_resampling_step = torch.clip( self.max_episode_length - self.episode_length_buf[env_ids] - (remaining_dist / (0.8 * self.max_lin_vel * self.dt + 1e-9)), min=0.0, max=self.cfg.commands.resampling_time / self.dt, ) zero_mask = (rand_prob >= min_prob) * (rand_prob < max_prob) * (next_resampling_step > 0.0) zero_env_ids = env_ids[zero_mask] if len(zero_env_ids) > 0: self.commands[zero_env_ids, :2] = 0.0 self.commands_resampling_step[zero_env_ids] = next_resampling_step[zero_mask] if self.cfg.commands.limit_ang_vel_at_zero_command_prob > 0.0: ang_vel_rand = torch.rand(len(zero_env_ids), device=self.device) # independent distribution add_ang_mask = ang_vel_rand < self.cfg.commands.limit_ang_vel_at_zero_command_prob add_ang_env_ids = zero_env_ids[add_ang_mask] if len(add_ang_env_ids) > 0: direction_rand = torch.rand(len(add_ang_env_ids), device=self.device) self.commands[add_ang_env_ids, 2] = torch.where( direction_rand < 0.5, self.env_command_ranges["ang_vel_yaw"][add_ang_env_ids, 0], self.env_command_ranges["ang_vel_yaw"][add_ang_env_ids, 1] ) if self.cfg.commands.heading_command: self.stop_heading[add_ang_env_ids] = True min_prob += self.zero_command_proba # turn over zero command time if self.cfg.init_state.turn_over and (self.turn_over_timer[env_ids] > 0).any(): zero_mask = self.turn_over_timer[env_ids] > 0 zero_env_ids = env_ids[zero_mask] self.commands[zero_env_ids, :3] = 0.0 self.stop_heading[zero_env_ids] = True self.commands_xy_accumulation[env_ids] += self.commands[env_ids, :2] def _compute_torques(self, actions): """ Compute torques from actions. Actions can be interpreted as position or velocity targets given to a PD controller, or directly as scaled torques. [NOTE]: torques must have the same dimension as the number of DOFs, even if some DOFs are not actuated. Args: actions (torch.Tensor): Actions Returns: [torch.Tensor]: Torques sent to the simulation """ #pd controller actions_scaled = actions * self.cfg.control.action_scale control_type = self.cfg.control.control_type p_gains = self.p_gains * self.p_gains_multiplier d_gains = self.d_gains * self.d_gains_multiplier if control_type=="P": torques = p_gains*(actions_scaled + self.default_dof_pos - self.dof_pos + self.motor_zero_offsets) - d_gains*self.dof_vel elif control_type=="V": torques = p_gains*(actions_scaled - self.dof_vel) - d_gains*(self.dof_vel - self.last_dof_vel)/self.sim_params.dt elif control_type=="T": torques = actions_scaled else: raise NameError(f"Unknown controller type: {control_type}") return torch.clip(torques, -self.torque_limits, self.torque_limits) def _reset_dofs(self, env_ids): """ Resets DOF position and velocities of selected environmments Positions are randomly selected within 0.5:1.5 x default positions. Velocities are set to zero. Args: env_ids (List[int]): Environemnt ids """ self.dof_pos[env_ids] = self.default_dof_pos * torch_rand_float(0.5, 1.5, (len(env_ids), self.num_dof), device=self.device) self.dof_vel[env_ids] = 0. env_ids_int32 = env_ids.to(dtype=torch.int32) self.gym.set_dof_state_tensor_indexed(self.sim, gymtorch.unwrap_tensor(self.dof_state), gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32)) def _reset_root_states(self, env_ids): """ Resets ROOT states position and velocities of selected environmments Sets base position based on the curriculum Selects randomized base velocities within -0.5:0.5 [m/s, rad/s] Args: env_ids (List[int]): Environemnt ids """ if self.cfg.init_state.turn_over: self.turn_over_timer[env_ids] = 0.0 # base position random_yaw = torch_rand_float(-np.pi, np.pi, (len(env_ids), 1), device=self.device).squeeze(1) def get_quat(target_yaws, roll: float): roll_tensor = torch.full((len(target_yaws),), roll, device=self.device) pitch_tensor = torch.zeros((len(target_yaws),), device=self.device) quat = quat_from_euler_xyz(roll_tensor, pitch_tensor, target_yaws) return quat base_init_state = self.base_init_state.reshape(1, -1).repeat(len(env_ids), 1) if self.cfg.init_state.turn_over: rand_prob = torch.rand(len(env_ids), device=self.device) proportions = self.cfg.init_state.turn_over_proportions init_heights = self.cfg.init_state.turn_over_init_heights min_prob, max_prob = 0.0, proportions[0] back_mask = (rand_prob >= min_prob) * (rand_prob < max_prob) # backflip if back_mask.any(): heights = torch_rand_float(init_heights['backflip'][0], init_heights['backflip'][1], (torch.sum(back_mask), 1), device=self.device).squeeze(1) base_init_state[back_mask, 2] = heights # z base_init_state[back_mask, 3:7] = get_quat(random_yaw[back_mask], np.pi) if self.cfg.init_state.turn_over: self.turn_over_timer[env_ids[back_mask]] = self.cfg.commands.turn_over_zero_time['backflip'] min_prob = max_prob max_prob += proportions[1] side_mask = (rand_prob >= min_prob) * (rand_prob < max_prob) # sideflip if side_mask.any(): side_ids = torch.nonzero(side_mask, as_tuple=False).flatten() heights = torch_rand_float(init_heights['sideflip'][0], init_heights['sideflip'][1], (len(side_ids), 1), device=self.device).squeeze(1) base_init_state[side_mask, 2] = heights # z side_rand_prob = torch.rand(len(side_ids), device=self.device) pos_side_mask = side_rand_prob < 0.5 neg_side_mask = ~pos_side_mask if pos_side_mask.any(): pos_ids = side_ids[pos_side_mask] base_init_state[pos_ids, 3:7] = get_quat(random_yaw[pos_ids], np.pi/2) if neg_side_mask.any(): neg_ids = side_ids[neg_side_mask] base_init_state[neg_ids, 3:7] = get_quat(random_yaw[neg_ids], -np.pi/2) if self.cfg.init_state.turn_over: self.turn_over_timer[env_ids[side_mask]] = self.cfg.commands.turn_over_zero_time['sideflip'] min_prob = max_prob max_prob += proportions[2] noflip_mask = (rand_prob >= min_prob) * (rand_prob < max_prob) # noflip if noflip_mask.any(): noflip_indices = torch.nonzero(noflip_mask, as_tuple=False).flatten() base_init_state[noflip_mask, 3:7] = get_quat(random_yaw[noflip_indices], 0.0) else: base_init_state[:, 3:7] = get_quat(random_yaw, 0.0) if self.custom_origins: self.root_states[env_ids] = base_init_state self.root_states[env_ids, :3] += self.env_origins[env_ids] self.root_states[env_ids, :2] += torch_rand_float(-1., 1., (len(env_ids), 2), device=self.device) # xy position within 1m of the center else: self.root_states[env_ids] = base_init_state self.root_states[env_ids, :3] += self.env_origins[env_ids] # base velocities self.root_states[env_ids, 7:13] = torch_rand_float(-0.5, 0.5, (len(env_ids), 6), device=self.device) # [7:10]: lin vel, [10:13]: ang vel env_ids_int32 = env_ids.to(dtype=torch.int32) self.gym.set_actor_root_state_tensor_indexed(self.sim, gymtorch.unwrap_tensor(self.root_states), gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32)) def _push_robots(self): """ Random pushes the robots. Emulates an impulse by setting a randomized base velocity. """ env_ids = torch.arange(self.num_envs, device=self.device) push_env_ids = env_ids[self.episode_length_buf[env_ids] % int(self.cfg.domain_rand.push_interval) == 0] if len(push_env_ids) == 0: return max_vel = self.cfg.domain_rand.max_push_vel_xy max_push_ang = self.cfg.domain_rand.max_push_ang_vel self.root_states[:, 7:9] = torch_rand_float(-max_vel, max_vel, (self.num_envs, 2), device=self.device) # lin vel x/y self.root_states[:, 10:13] = torch_rand_float(-max_push_ang, max_push_ang, (self.num_envs, 3), device=self.device) # ang vel x/y/z env_ids_int32 = push_env_ids.to(dtype=torch.int32) self.gym.set_actor_root_state_tensor_indexed(self.sim, gymtorch.unwrap_tensor(self.root_states), gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32)) def update_command_curriculum(self, env_ids): """ Implements a curriculum of increasing commands Args: env_ids (List[int]): ids of environments being reset """ # If the tracking reward is above 80% of the maximum, increase the range of commands if torch.mean(self.episode_sums["tracking_lin_vel"][env_ids]) / self.max_episode_length > 0.8 * self.reward_scales["tracking_lin_vel"]: self.command_ranges["lin_vel_x"][0] = np.clip(self.command_ranges["lin_vel_x"][0] - 0.5, -self.cfg.commands.max_curriculum, 0.) self.command_ranges["lin_vel_x"][1] = np.clip(self.command_ranges["lin_vel_x"][1] + 0.5, 0., self.cfg.commands.max_curriculum) def _get_noise_scale_vec(self, cfg): """ Sets a vector used to scale the noise added to the observations. [NOTE]: Must be adapted when changing the observations structure Args: cfg (Dict): Environment config file Returns: [torch.Tensor]: Vector of scales used to multiply a uniform distribution in [-1, 1] """ 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.lin_vel * noise_level * self.obs_scales.lin_vel noise_vec[3:6] = noise_scales.ang_vel * noise_level * self.obs_scales.ang_vel noise_vec[6:9] = noise_scales.gravity * noise_level noise_vec[9:12] = 0. # commands noise_vec[12:12+self.num_actions] = noise_scales.dof_pos * noise_level * self.obs_scales.dof_pos noise_vec[12+self.num_actions:12+2*self.num_actions] = noise_scales.dof_vel * noise_level * self.obs_scales.dof_vel noise_vec[12+2*self.num_actions:12+3*self.num_actions] = 0. # previous actions return noise_vec #---------------------------------------- def _init_buffers(self): """ Initialize torch tensors which will contain simulation states and processed quantities """ # get gym GPU state tensors actor_root_state = self.gym.acquire_actor_root_state_tensor(self.sim) dof_state_tensor = self.gym.acquire_dof_state_tensor(self.sim) net_contact_forces = self.gym.acquire_net_contact_force_tensor(self.sim) rigid_body_state = self.gym.acquire_rigid_body_state_tensor(self.sim) self.gym.refresh_dof_state_tensor(self.sim) self.gym.refresh_actor_root_state_tensor(self.sim) self.gym.refresh_net_contact_force_tensor(self.sim) self.gym.refresh_rigid_body_state_tensor(self.sim) # create some wrapper tensors for different slices self.root_states = gymtorch.wrap_tensor(actor_root_state) self.dof_state = gymtorch.wrap_tensor(dof_state_tensor) self.dof_pos = self.dof_state.view(self.num_envs, self.num_dof, 2)[..., 0] self.dof_vel = self.dof_state.view(self.num_envs, self.num_dof, 2)[..., 1] self.base_quat = self.root_states[:, 3:7] self.rpy = get_euler_xyz_in_tensor(self.base_quat) self.base_pos = self.root_states[:self.num_envs, 0:3] self.contact_forces = gymtorch.wrap_tensor(net_contact_forces).view(self.num_envs, -1, 3) # shape: num_envs, num_bodies, xyz axis self.rigid_body_states = gymtorch.wrap_tensor(rigid_body_state) if self.cfg.terrain.measure_heights: self.height_points = self._init_height_points() x_points = self.height_points[0, :, 0] y_points = self.height_points[0, :, 1] x_mask = (x_points >= -0.2) & (x_points <= 0.2) # 0.4m length y_mask = (y_points >= -0.15) & (y_points <= 0.15) # 0.3m width self.base_height_scan_mask = (x_mask & y_mask).float() self.num_base_height_scan_points = self.base_height_scan_mask.sum() assert self.num_base_height_scan_points > 0, "No height scan points within the specified area." self.measured_heights = 0 # initialize some data used later on self.common_step_counter = 0 self.extras = {} self.noise_scale_vec = self._get_noise_scale_vec(self.cfg) self.gravity_vec = to_torch(get_axis_params(-1., self.up_axis_idx), device=self.device).repeat((self.num_envs, 1)) self.forward_vec = to_torch([1., 0., 0.], device=self.device).repeat((self.num_envs, 1)) self.torques = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False) self.p_gains = torch.zeros(self.num_actions, dtype=torch.float, device=self.device, requires_grad=False) self.d_gains = torch.zeros(self.num_actions, dtype=torch.float, device=self.device, requires_grad=False) self.actions = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False) self.last_actions = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False) self.last_dof_vel = torch.zeros_like(self.dof_vel) self.last_root_vel = torch.zeros_like(self.root_states[:, 7:13]) self.commands = torch.zeros(self.num_envs, self.cfg.commands.num_commands, dtype=torch.float, device=self.device, requires_grad=False) # x vel, y vel, yaw vel, heading self.commands_scale = torch.tensor([self.obs_scales.lin_vel, self.obs_scales.lin_vel, self.obs_scales.ang_vel], device=self.device, requires_grad=False,) # TODO change this self.commands_resampling_step = torch.zeros(self.num_envs, dtype=torch.float, device=self.device, requires_grad=False) self.commands_xy_accumulation = torch.zeros(self.num_envs, 2, dtype=torch.float, device=self.device, requires_grad=False) self.zero_command_proba = 0.0 self.feet_air_time = torch.zeros(self.num_envs, self.feet_indices.shape[0], dtype=torch.float, device=self.device, requires_grad=False) self.last_contacts = torch.zeros(self.num_envs, len(self.feet_indices), dtype=torch.bool, device=self.device, requires_grad=False) self.base_lin_vel = quat_rotate_inverse(self.base_quat, self.root_states[:, 7:10]) self.base_ang_vel = quat_rotate_inverse(self.base_quat, self.root_states[:, 10:13]) self.projected_gravity = quat_rotate_inverse(self.base_quat, self.gravity_vec) self.max_move_distance = torch.zeros(self.num_envs, dtype=torch.float, device=self.device, requires_grad=False) self.stop_heading = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device, requires_grad=False) self.last_is_limit_vel = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device, requires_grad=False) self.motor_strengths = torch.ones(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False) self.limit_vel_prob = self.cfg.commands.limit_vel_prob self.limit_vel_comb = torch.tensor(list(product( self.cfg.commands.limit_vel["lin_vel_x"], self.cfg.commands.limit_vel["lin_vel_y"], self.cfg.commands.limit_vel["ang_vel_yaw"] )), device=self.device, requires_grad=False) self.last_robot_props_update_step = torch.zeros(self.num_envs, dtype=torch.long, device=self.device, requires_grad=False) self.turn_over_timer = torch.zeros(self.num_envs, dtype=torch.float, device=self.device, requires_grad=False) self.env_command_ranges = { 'lin_vel_x': torch.zeros((self.num_envs, 2), dtype=torch.float, device=self.device, requires_grad=False), 'lin_vel_y': torch.zeros((self.num_envs, 2), dtype=torch.float, device=self.device, requires_grad=False), 'ang_vel_yaw': torch.zeros((self.num_envs, 2), dtype=torch.float, device=self.device, requires_grad=False), 'heading': torch.zeros((self.num_envs, 2), dtype=torch.float, device=self.device, requires_grad=False), } self._update_env_command_ranges() # joint positions offsets and PD gains self.default_dof_pos = torch.zeros(self.num_dof, dtype=torch.float, device=self.device, requires_grad=False) for i in range(self.num_dofs): name = self.dof_names[i] angle = self.cfg.init_state.default_joint_angles[name] self.default_dof_pos[i] = angle found = False for dof_name in self.cfg.control.stiffness.keys(): if dof_name in name: self.p_gains[i] = self.cfg.control.stiffness[dof_name] self.d_gains[i] = self.cfg.control.damping[dof_name] found = True if not found: self.p_gains[i] = 0. self.d_gains[i] = 0. if self.cfg.control.control_type in ["P", "V"]: print(f"PD gain of joint {name} were not defined, setting them to zero") self.default_dof_pos = self.default_dof_pos.unsqueeze(0) def _update_env_command_ranges(self): """ Update environment-wise command ranges based on current command ranges and terrain type """ if not hasattr(self, 'terrain_ids'): return for terrain_id, terrain_command_ranges in enumerate(self.cfg.commands.terrain_max_command_ranges): env_ids = (self.terrain_ids == terrain_id).nonzero(as_tuple=False).flatten() if len(env_ids) == 0: continue self.env_command_ranges['lin_vel_x'][env_ids, 0] = max( terrain_command_ranges['lin_vel_x'][0], self.command_ranges['lin_vel_x'][0], ) self.env_command_ranges['lin_vel_x'][env_ids, 1] = min( terrain_command_ranges['lin_vel_x'][1], self.command_ranges['lin_vel_x'][1] ) self.env_command_ranges['lin_vel_y'][env_ids, 0] = max( terrain_command_ranges['lin_vel_y'][0], self.command_ranges['lin_vel_y'][0] ) self.env_command_ranges['lin_vel_y'][env_ids, 1] = min( terrain_command_ranges['lin_vel_y'][1], self.command_ranges['lin_vel_y'][1] ) self.env_command_ranges['ang_vel_yaw'][env_ids, 0] = max( terrain_command_ranges['ang_vel_yaw'][0], self.command_ranges['ang_vel_yaw'][0] ) self.env_command_ranges['ang_vel_yaw'][env_ids, 1] = min( terrain_command_ranges['ang_vel_yaw'][1], self.command_ranges['ang_vel_yaw'][1] ) if self.cfg.commands.heading_command: self.env_command_ranges['heading'][env_ids, 0] = max( terrain_command_ranges['heading'][0], self.command_ranges['heading'][0] ) self.env_command_ranges['heading'][env_ids, 1] = min( terrain_command_ranges['heading'][1], self.command_ranges['heading'][1] ) def _prepare_reward_function(self): """ Prepares a list of reward functions, whcih will be called to compute the total reward. Looks for self._reward_, where are names of all non zero reward scales in the cfg. """ # remove zero scales + multiply non-zero ones by dt def update_scales(scales): for key in list(scales.keys()): scale = scales[key] if scale==0: scales.pop(key) else: scales[key] *= self.dt update_scales(self.reward_scales) if self.cfg.init_state.turn_over: update_scales(self.reward_turn_over_scales) # prepare list of functions self.reward_functions = [] self.reward_names = [] names = set() names.update(list(self.reward_scales.keys())) if self.cfg.init_state.turn_over: names.update(list(self.reward_turn_over_scales.keys())) for name in names: if name=="termination": continue self.reward_names.append(name) name = '_reward_' + name self.reward_functions.append(getattr(self, name)) # reward episode sums self.episode_sums = {name: torch.zeros(self.num_envs, dtype=torch.float, device=self.device, requires_grad=False) for name in names} def _create_ground_plane(self): """ Adds a ground plane to the simulation, sets friction and restitution based on the cfg. """ plane_params = gymapi.PlaneParams() plane_params.normal = gymapi.Vec3(0.0, 0.0, 1.0) plane_params.static_friction = self.cfg.terrain.static_friction plane_params.dynamic_friction = self.cfg.terrain.dynamic_friction plane_params.restitution = self.cfg.terrain.restitution self.gym.add_ground(self.sim, plane_params) def _create_envs(self): """ Creates environments: 1. loads the robot URDF/MJCF asset, 2. For each environment 2.1 creates the environment, 2.2 calls DOF and Rigid shape properties callbacks, 2.3 create actor with these properties and add them to the env 3. Store indices of different bodies of the robot """ asset_path = self.cfg.asset.file.format(LEGGED_GYM_ROOT_DIR=LEGGED_GYM_ROOT_DIR) asset_root = os.path.dirname(asset_path) asset_file = os.path.basename(asset_path) asset_options = gymapi.AssetOptions() asset_options.default_dof_drive_mode = self.cfg.asset.default_dof_drive_mode asset_options.collapse_fixed_joints = self.cfg.asset.collapse_fixed_joints asset_options.replace_cylinder_with_capsule = self.cfg.asset.replace_cylinder_with_capsule asset_options.flip_visual_attachments = self.cfg.asset.flip_visual_attachments asset_options.fix_base_link = self.cfg.asset.fix_base_link asset_options.density = self.cfg.asset.density asset_options.angular_damping = self.cfg.asset.angular_damping asset_options.linear_damping = self.cfg.asset.linear_damping asset_options.max_angular_velocity = self.cfg.asset.max_angular_velocity asset_options.max_linear_velocity = self.cfg.asset.max_linear_velocity asset_options.armature = self.cfg.asset.armature asset_options.thickness = self.cfg.asset.thickness asset_options.disable_gravity = self.cfg.asset.disable_gravity self.robot_asset = self.gym.load_asset(self.sim, asset_root, asset_file, asset_options) self.num_dof = self.gym.get_asset_dof_count(self.robot_asset) self.num_bodies = self.gym.get_asset_rigid_body_count(self.robot_asset) dof_props_asset = self.gym.get_asset_dof_properties(self.robot_asset) rigid_shape_props_asset = self.gym.get_asset_rigid_shape_properties(self.robot_asset) # save body names from the asset body_names = self.gym.get_asset_rigid_body_names(self.robot_asset) self.dof_names = self.gym.get_asset_dof_names(self.robot_asset) self.num_bodies = len(body_names) self.num_dofs = len(self.dof_names) feet_names = [s for s in body_names if self.cfg.asset.foot_name in s] hip_names = [s for s in self.dof_names if 'hip' in s] penalized_contact_names = [] for name in self.cfg.asset.penalize_contacts_on: penalized_contact_names.extend([s for s in body_names if name in s]) termination_contact_names = [] for name in self.cfg.asset.terminate_after_contacts_on: termination_contact_names.extend([s for s in body_names if name in s]) base_init_state_list = self.cfg.init_state.pos + self.cfg.init_state.rot + self.cfg.init_state.lin_vel + self.cfg.init_state.ang_vel self.base_init_state = to_torch(base_init_state_list, device=self.device, requires_grad=False) start_pose = gymapi.Transform() start_pose.p = gymapi.Vec3(*self.base_init_state[:3]) # domain rand self.motor_zero_offsets = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False) self.p_gains_multiplier = torch.ones(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False) self.d_gains_multiplier = torch.ones(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False) if self.cfg.rewards.dynamic_sigma: self.dynamic_sigma_cfg = self.cfg.rewards.dynamic_sigma self.terrain_max_sigmas = torch.tensor(self.dynamic_sigma_cfg["max_sigma"], device=self.device, requires_grad=False) self._get_env_origins() env_lower = gymapi.Vec3(0., 0., 0.) env_upper = gymapi.Vec3(0., 0., 0.) self.actor_handles = [] self.envs = [] for i in range(self.num_envs): # create env instance env_handle = self.gym.create_env(self.sim, env_lower, env_upper, int(np.sqrt(self.num_envs))) pos = self.env_origins[i].clone() pos[:2] += torch_rand_float(-1., 1., (2,1), device=self.device).squeeze(1) start_pose.p = gymapi.Vec3(*pos) rigid_shape_props = self._process_rigid_shape_props(rigid_shape_props_asset, i) self.gym.set_asset_rigid_shape_properties(self.robot_asset, rigid_shape_props) actor_handle = self.gym.create_actor(env_handle, self.robot_asset, start_pose, self.cfg.asset.name, i, self.cfg.asset.self_collisions, 0) dof_props = self._process_dof_props(dof_props_asset, i) self.gym.set_actor_dof_properties(env_handle, actor_handle, dof_props) body_props = self.gym.get_actor_rigid_body_properties(env_handle, actor_handle) if i == 0: self.default_body_props = self.gym.get_actor_rigid_body_properties(env_handle, actor_handle) body_props = self._process_rigid_body_props(body_props, i) self.gym.set_actor_rigid_body_properties(env_handle, actor_handle, body_props, recomputeInertia=True) self.envs.append(env_handle) self.actor_handles.append(actor_handle) self.feet_indices = torch.zeros(len(feet_names), dtype=torch.long, device=self.device, requires_grad=False) for i in range(len(feet_names)): self.feet_indices[i] = self.gym.find_actor_rigid_body_handle(self.envs[0], self.actor_handles[0], feet_names[i]) self.hip_indices = torch.zeros(len(hip_names), dtype=torch.long, device=self.device, requires_grad=False) for i in range(len(hip_names)): self.hip_indices[i] = self.gym.find_actor_dof_handle(self.envs[0], self.actor_handles[0], hip_names[i]) self.penalised_contact_indices = torch.zeros(len(penalized_contact_names), dtype=torch.long, device=self.device, requires_grad=False) for i in range(len(penalized_contact_names)): self.penalised_contact_indices[i] = self.gym.find_actor_rigid_body_handle(self.envs[0], self.actor_handles[0], penalized_contact_names[i]) self.termination_contact_indices = torch.zeros(len(termination_contact_names), dtype=torch.long, device=self.device, requires_grad=False) for i in range(len(termination_contact_names)): self.termination_contact_indices[i] = self.gym.find_actor_rigid_body_handle(self.envs[0], self.actor_handles[0], termination_contact_names[i]) def _get_env_origins(self): """ Sets environment origins. On rough terrain the origins are defined by the terrain platforms. Otherwise create a grid. """ if self.cfg.terrain.mesh_type in ["heightfield", "trimesh"]: self.custom_origins = True self.env_origins = torch.zeros(self.num_envs, 3, device=self.device, requires_grad=False) # put robots at the origins defined by the terrain max_init_level = self.cfg.terrain.max_init_terrain_level if not self.cfg.terrain.curriculum: max_init_level = self.cfg.terrain.num_rows - 1 # random choice terrain levels and types for each env # self.terrain_levels = torch.randint(0, max_init_level+1, (self.num_envs,), device=self.device) # self.terrain_types = torch.randint(0, self.cfg.terrain.num_cols, (self.num_envs,), device=self.device) # levels and types in a round robin manner self.terrain_levels = torch.fmod(torch.arange(self.num_envs, device=self.device), max_init_level + 1) self.terrain_types = torch.div(torch.arange(self.num_envs, device=self.device), (self.num_envs / self.cfg.terrain.num_cols), rounding_mode="floor").to(torch.long) self.terrain_cols2id = torch.tensor(self.terrain.cols2id, device=self.device) if len(self.terrain_cols2id): self.terrain_ids = self.terrain_cols2id[self.terrain_types] self.max_terrain_level = self.cfg.terrain.num_rows self.terrain_origins = torch.from_numpy(self.terrain.env_origins).to(self.device).to(torch.float) self.env_origins[:] = self.terrain_origins[self.terrain_levels, self.terrain_types] else: self.custom_origins = False self.env_origins = torch.zeros(self.num_envs, 3, device=self.device, requires_grad=False) # create a grid of robots num_cols = np.floor(np.sqrt(self.num_envs)) num_rows = np.ceil(self.num_envs / num_cols) xx, yy = torch.meshgrid(torch.arange(num_rows), torch.arange(num_cols)) spacing = self.cfg.env.env_spacing self.env_origins[:, 0] = spacing * xx.flatten()[:self.num_envs] self.env_origins[:, 1] = spacing * yy.flatten()[:self.num_envs] self.env_origins[:, 2] = 0. def _parse_cfg(self, cfg): self.dt = self.cfg.control.decimation * self.sim_params.dt self.obs_scales = self.cfg.normalization.obs_scales self.reward_scales = class_to_dict(self.cfg.rewards.scales) self.reward_turn_over_scales = class_to_dict(self.cfg.rewards.turn_over_scales) self.command_ranges = class_to_dict(self.cfg.commands.ranges) self.max_lin_vel = max(abs(self.command_ranges["lin_vel_x"][0]), abs(self.command_ranges["lin_vel_x"][1]), abs(self.command_ranges["lin_vel_y"][0]), abs(self.command_ranges["lin_vel_y"][1])) self.cfg.commands.command_range_curriculum = sorted(self.cfg.commands.command_range_curriculum, key=lambda x: x['iter'], reverse=True) self.max_episode_length_s = self.cfg.env.episode_length_s self.max_episode_length = np.ceil(self.max_episode_length_s / self.dt) self.cfg.domain_rand.push_interval = np.ceil(self.cfg.domain_rand.push_interval_s / self.dt) def _create_heightfield(self): """ Adds a heightfield terrain to the simulation, sets parameters based on the cfg. """ hf_params = gymapi.HeightFieldParams() hf_params.column_scale = self.terrain.cfg.horizontal_scale hf_params.row_scale = self.terrain.cfg.horizontal_scale hf_params.vertical_scale = self.terrain.cfg.vertical_scale hf_params.nbRows = self.terrain.tot_cols hf_params.nbColumns = self.terrain.tot_rows hf_params.transform.p.x = -self.terrain.cfg.border_size hf_params.transform.p.y = -self.terrain.cfg.border_size hf_params.transform.p.z = 0.0 hf_params.static_friction = self.cfg.terrain.static_friction hf_params.dynamic_friction = self.cfg.terrain.dynamic_friction hf_params.restitution = self.cfg.terrain.restitution self.gym.add_heightfield(self.sim, self.terrain.heightsamples, hf_params) self.height_samples = torch.tensor(self.terrain.heightsamples).view(self.terrain.tot_rows, self.terrain.tot_cols).to(self.device) def _create_trimesh(self): """ Adds a triangle mesh terrain to the simulation, sets parameters based on the cfg. # """ tm_params = gymapi.TriangleMeshParams() tm_params.nb_vertices = self.terrain.vertices.shape[0] tm_params.nb_triangles = self.terrain.triangles.shape[0] tm_params.transform.p.x = -self.terrain.cfg.border_size tm_params.transform.p.y = -self.terrain.cfg.border_size tm_params.transform.p.z = 0.0 tm_params.static_friction = self.cfg.terrain.static_friction tm_params.dynamic_friction = self.cfg.terrain.dynamic_friction tm_params.restitution = self.cfg.terrain.restitution self.gym.add_triangle_mesh(self.sim, self.terrain.vertices.flatten(order='C'), self.terrain.triangles.flatten(order='C'), tm_params) self.height_samples = torch.tensor(self.terrain.heightsamples).view(self.terrain.tot_rows, self.terrain.tot_cols).to(self.device) def _update_terrain_curriculum(self, env_ids): """ Implements the game-inspired curriculum. Args: env_ids (List[int]): ids of environments being reset """ # Implement Terrain curriculum if not self.init_done or self.cfg.terrain.mesh_type == 'plane': # don't change on initial reset return # distance = torch.norm(self.root_states[env_ids, :2] - self.env_origins[env_ids, :2], dim=1) distance = self.max_move_distance[env_ids] # robots that walked far enough progress to harder terains move_up = distance > self.terrain.env_length / 2 if self.cfg.terrain.move_down_by_acuumulated_xy_command: move_down = (distance < torch.norm(self.commands_xy_accumulation[env_ids], dim=1) * (self.cfg.commands.resampling_time * (1 - self.zero_command_proba)) * 0.5) * ~move_up else: # robots that walked less than half of their required distance go to simpler terrains move_down = (distance < torch.norm(self.commands[env_ids, :2], dim=1) * self.max_episode_length_s * 0.5) * ~move_up self.terrain_levels[env_ids] += 1 * move_up - 1 * move_down # Robots that solve the last level are sent to a random one self.terrain_levels[env_ids] = torch.where(self.terrain_levels[env_ids]>=self.max_terrain_level, torch.randint_like(self.terrain_levels[env_ids], self.max_terrain_level), torch.clip(self.terrain_levels[env_ids], 0)) # (the minumum level is zero) self.env_origins[env_ids] = self.terrain_origins[self.terrain_levels[env_ids], self.terrain_types[env_ids]] self.max_move_distance[env_ids] = 0.0 def _init_height_points(self): """ Returns points at which the height measurments are sampled (in base frame) Returns: [torch.Tensor]: Tensor of shape (num_envs, self.num_height_points, 3) """ y = torch.tensor(self.cfg.terrain.measured_points_y, device=self.device, requires_grad=False) x = torch.tensor(self.cfg.terrain.measured_points_x, device=self.device, requires_grad=False) grid_x, grid_y = torch.meshgrid(x, y) self.num_height_points = grid_x.numel() points = torch.zeros(self.num_envs, self.num_height_points, 3, device=self.device, requires_grad=False) points[:, :, 0] = grid_x.flatten() points[:, :, 1] = grid_y.flatten() return points def _get_heights(self, env_ids=None): """ Samples heights of the terrain at required points around each robot. The points are offset by the base's position and rotated by the base's yaw Args: env_ids (List[int], optional): Subset of environments for which to return the heights. Defaults to None. Raises: NameError: [description] Returns: [type]: [description] """ if self.cfg.terrain.mesh_type == "plane": return torch.zeros(self.num_envs, self.num_height_points, device=self.device, requires_grad=False) elif self.cfg.terrain.mesh_type == "none": raise NameError("Can't measure height with terrain mesh type 'none'") if env_ids: points = quat_apply_yaw(self.base_quat[env_ids].repeat(1, self.num_height_points), self.height_points[env_ids]) + (self.root_states[env_ids, :3]).unsqueeze(1) else: points = quat_apply_yaw(self.base_quat.repeat(1, self.num_height_points), self.height_points) + (self.root_states[:, :3]).unsqueeze(1) points += self.terrain.cfg.border_size points = (points / self.terrain.cfg.horizontal_scale).long() px = points[:, :, 0].view(-1) py = points[:, :, 1].view(-1) px = torch.clip(px, 0, self.height_samples.shape[0] - 2) py = torch.clip(py, 0, self.height_samples.shape[1] - 2) heights1 = self.height_samples[px, py] heights2 = self.height_samples[px + 1, py] heights3 = self.height_samples[px, py + 1] heights = torch.min(heights1, heights2) heights = torch.min(heights, heights3) return heights.view(self.num_envs, -1) * self.terrain.cfg.vertical_scale #------------ reward functions---------------- def _reward_lin_vel_z(self): # Penalize z axis base linear velocity return torch.square(self.base_lin_vel[:, 2]) def _reward_ang_vel_xy(self): # Penalize xy axes base angular velocity return torch.sum(torch.square(self.base_ang_vel[:, :2]), dim=1) def _reward_orientation(self): # Penalize non flat base orientation return torch.sum(torch.square(self.projected_gravity[:, :2]), dim=1) # def _reward_base_height(self): # # Penalize base height away from target # base_height = self.root_states[:, 2] # return torch.square(base_height - self.cfg.rewards.base_height_target) def _reward_base_height(self): # Penalize base height away from target contact = self.contact_forces[:, self.feet_indices, 2] > 1. if not hasattr(self, 'last_contacts2'): self.last_contacts2 = torch.zeros_like(contact) contact_filt = torch.logical_or(contact, self.last_contacts2) # (N, 4) self.last_contacts2 = contact feet_pos = self.rigid_body_states.view(self.num_envs, self.num_bodies, 13)[:, self.feet_indices, 0:3] num_feet_contact = torch.sum(contact_filt, dim=1, keepdim=True).clamp(min=1.0) # (N, 1) feet_contact_pos = (feet_pos * contact_filt.unsqueeze(-1)).sum(dim=1) / num_feet_contact # (N, 3) base_pos = self.root_states[:, 0:3] delta_pos = feet_contact_pos - base_pos base_height = (delta_pos * self.projected_gravity).sum(1) # (N,) rew = torch.square(base_height - self.cfg.rewards.base_height_target) * (contact_filt.sum(1) > 0) return rew def _reward_torques(self): # Penalize torques return torch.sum(torch.square(self.torques), dim=1) def _reward_dof_vel(self): # Penalize dof velocities return torch.sum(torch.square(self.dof_vel), dim=1) def _reward_dof_acc(self): # Penalize dof accelerations return torch.sum(torch.square((self.last_dof_vel - self.dof_vel) / self.dt), dim=1) def _reward_action_rate(self): # Penalize changes in actions return torch.sum(torch.square(self.last_actions - self.actions), dim=1) def _reward_collision(self): # Penalize collisions on selected bodies return torch.sum(1.*(torch.norm(self.contact_forces[:, self.penalised_contact_indices, :], dim=-1) > 0.1), dim=1) def _reward_termination(self): # Terminal reward / penalty return self.reset_buf * ~self.time_out_buf def _reward_dof_pos_limits(self): # Penalize dof positions too close to the limit out_of_limits = -(self.dof_pos - self.dof_pos_limits[:, 0]).clip(max=0.) # lower limit out_of_limits += (self.dof_pos - self.dof_pos_limits[:, 1]).clip(min=0.) return torch.sum(out_of_limits, dim=1) def _reward_dof_vel_limits(self): # Penalize dof velocities too close to the limit # clip to max error = 1 rad/s per joint to avoid huge penalties return torch.sum((torch.abs(self.dof_vel) - self.dof_vel_limits*self.cfg.rewards.soft_dof_vel_limit).clip(min=0., max=1.), dim=1) def _reward_torque_limits(self): # penalize torques too close to the limit return torch.sum((torch.abs(self.torques) - self.torque_limits*self.cfg.rewards.soft_torque_limit).clip(min=0.), dim=1) def _get_dynamic_sigma(self, target_vel_abs, v_min, v_max): # compute dynamic sigma based on terrain level default_sigma = self.cfg.rewards.tracking_sigma if not self.cfg.terrain.curriculum or self.cfg.rewards.dynamic_sigma is None or not hasattr(self, 'terrain_ids'): return torch.full_like(target_vel_abs, default_sigma) target_sigmas = self.terrain_max_sigmas[self.terrain_ids] sigma = torch.full_like(target_vel_abs, default_sigma) # based on velocity ranges, compute sigma # v_min <= v < v_max (linear interpolation) mask = (target_vel_abs >= v_min) & (target_vel_abs < v_max) if mask.any(): ratio = (target_vel_abs[mask] - v_min) / (v_max - v_min) sigma[mask] = default_sigma + ratio * (target_sigmas[mask] - default_sigma) # v >= v_max mask = target_vel_abs >= v_max if mask.any(): sigma[mask] = target_sigmas[mask] # based on terrain level, compute sigma level_scale = torch.clamp(torch.exp((self.terrain_levels.float() + 1.0) / 10.0) - 1.0, max=1.0) sigma = default_sigma + level_scale * (sigma - default_sigma) return sigma def _reward_tracking_lin_vel(self): # Tracking of linear velocity commands (xy axes) if self.cfg.rewards.dynamic_sigma is None: sigma_x = sigma_y = self.cfg.rewards.tracking_sigma else: vmin = self.dynamic_sigma_cfg["min_lin_vel"] vmax = self.dynamic_sigma_cfg["max_lin_vel"] sigma_x = self._get_dynamic_sigma(torch.abs(self.commands[:, 0]), vmin, vmax) sigma_y = self._get_dynamic_sigma(torch.abs(self.commands[:, 1]), vmin, vmax) lin_vel_error_sq = torch.square(self.commands[:, :2] - self.base_lin_vel[:, :2]) scaled_error = lin_vel_error_sq[:, 0] / sigma_x + lin_vel_error_sq[:, 1] / sigma_y return torch.exp(-scaled_error) def _reward_tracking_ang_vel(self): # Tracking of angular velocity commands (yaw) if self.cfg.rewards.dynamic_sigma is None: sigma = self.cfg.rewards.tracking_sigma else: vmin = self.dynamic_sigma_cfg["min_ang_vel"] vmax = self.dynamic_sigma_cfg["max_ang_vel"] sigma = self._get_dynamic_sigma(torch.abs(self.commands[:, 2]), vmin, vmax) ang_vel_error_sq = torch.square(self.commands[:, 2] - self.base_ang_vel[:, 2]) return torch.exp(-ang_vel_error_sq/sigma) def _reward_feet_air_time(self): # Reward long steps # Need to filter the contacts because the contact reporting of PhysX is unreliable on meshes contact = self.contact_forces[:, self.feet_indices, 2] > 1. contact_filt = torch.logical_or(contact, self.last_contacts) self.last_contacts = contact first_contact = (self.feet_air_time > 0.) * contact_filt self.feet_air_time += self.dt rew_airTime = torch.sum((self.feet_air_time - 0.5) * first_contact, dim=1) # reward only on first contact with the ground rew_airTime *= torch.norm(self.commands[:, :2], dim=1) > 0.1 #no reward for zero command self.feet_air_time *= ~contact_filt return rew_airTime def _reward_stumble(self): # Penalize feet hitting vertical surfaces return torch.any(torch.norm(self.contact_forces[:, self.feet_indices, :2], dim=2) >\ 5 *torch.abs(self.contact_forces[:, self.feet_indices, 2]), dim=1) def _reward_stand_still(self): # Penalize motion at zero commands return torch.sum(torch.abs(self.dof_pos - self.default_dof_pos), dim=1) * (torch.norm(self.commands[:, :2], dim=1) < 0.1) def _reward_feet_contact_forces(self): # penalize high contact forces return torch.sum((torch.norm(self.contact_forces[:, self.feet_indices, :], dim=-1) - self.cfg.rewards.max_contact_force).clip(min=0.), dim=1) def _reward_action_smoothness(self): # a_t - 2a_{t-1} + a_{t-2} if not hasattr(self, 'last_last_actions'): self.last_last_actions = torch.zeros_like(self.last_actions) rew = torch.sum((self.actions - 2 * self.last_actions + self.last_last_actions).pow(2), dim=1) self.last_last_actions[:] = self.last_actions[:] return rew def _reward_dof_power(self): # Penalize power consumption power = self.torques * self.dof_vel rew = torch.sum(torch.abs(power), dim=1) return rew def _get_base_height(self): if not self.cfg.terrain.measure_heights: return self.root_states[:, 2] # 根据高度扫描点计算base link到地面估计高度 masked_heights = self.measured_heights * self.base_height_scan_mask.unsqueeze(0) sum_heights = masked_heights.sum(dim=1) estimated_ground_z = sum_heights / self.num_base_height_scan_points base_z = self.root_states[:, 2] base_height = base_z - estimated_ground_z # (N,) return base_height def _reward_correct_base_height(self): base_height = self._get_base_height() rew = torch.square(base_height - self.cfg.rewards.base_height_target) return rew def _reward_feet_regulation(self): # CTS抬腿正则奖励, 在脚末端速度增大同时, 要求高度尽可能高 base_height = self._get_base_height() # 更新刚体空间位置 (开悟比赛无法修改环境, 只能在奖励中完成计算了) feet_pos = self.rigid_body_states.view(self.num_envs, self.num_bodies, 13)[:, self.feet_indices, 0:3] feet_xy_vel = self.rigid_body_states.view(self.num_envs, self.num_bodies, 13)[:, self.feet_indices, 7:9] base_pos = self.root_states[:, 0:3].unsqueeze(1) delta_feet = feet_pos - base_pos feet2base_height = (delta_feet * self.projected_gravity.unsqueeze(1)).sum(-1) # 脚相对于身体的高度 (N, 4) feet_height = torch.clamp(base_height.unsqueeze(1) - feet2base_height, min=0.0) # 脚相对于地面的高度 (N, 4) rew = (feet_xy_vel.pow(2).sum(-1) * torch.exp(-feet_height / (0.025 * self.cfg.rewards.base_height_target))).sum(-1) return rew def _reward_similar_to_default(self): # Penalize joint poses far away from default pose return torch.sum(torch.abs(self.dof_pos - self.default_dof_pos), dim=1) def _reward_hip_to_default(self): hip_dof_names = ['FL_hip_joint', 'FR_hip_joint', 'RL_hip_joint', 'RR_hip_joint'] hip_dof_indices = [0, 3, 6, 9] hip_pos = self.dof_pos[:, hip_dof_indices] default_hip_pos = self.default_dof_pos[:, hip_dof_indices] return torch.sum(torch.abs(hip_pos - default_hip_pos), dim=1) def _reward_upright(self): return (-1 - self.projected_gravity[:, 2]) / 2 def _reward_legs_distance(self): # Penalize legs being too close to each other # feet_names: [FL_foot, FR_foot, RL_foot, RR_foot] feet_pos_world = self.rigid_body_states.view(self.num_envs, self.num_bodies, 13)[:, self.feet_indices, 0:3] # (N, 4, 3) base_pos = self.root_states[:, 0:3] # (N, 3) base_quat = self.base_quat # (N, 4) feet_pos_relative_world = feet_pos_world - base_pos.unsqueeze(1) # (N, 4, 3) local_pos = quat_rotate_inverse( base_quat.repeat_interleave(4, dim=0), feet_pos_relative_world.reshape(-1, 3) ).reshape(self.num_envs, 4, 3) # (N, 4, 3) dy_front = local_pos[:, 0, 1] - local_pos[:, 1, 1] # (N,) dy_rear = local_pos[:, 2, 1] - local_pos[:, 3, 1] # (N,) min_dist = self.cfg.rewards.min_legs_distance rew_front = torch.square(torch.clamp(min_dist - dy_front, min=0.0)) rew_rear = torch.square(torch.clamp(min_dist - dy_rear, min=0.0)) return rew_front + rew_rear