# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause # # 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. # # Copyright (c) 2021 ETH Zurich, Nikita Rudin from legged_gym import LEGGED_GYM_ROOT_DIR from collections import OrderedDict import os import json import time import numpy as np np.float = np.float32 import isaacgym from isaacgym import gymtorch, gymapi from isaacgym.torch_utils import * from legged_gym.envs import * from legged_gym.utils import get_args, export_policy_as_jit, task_registry, Logger from legged_gym.utils.helpers import update_class_from_dict from legged_gym.utils.observation import get_obs_slice from legged_gym.debugger import break_into_debugger import numpy as np import torch def create_recording_camera(gym, env_handle, resolution= (1920, 1080), h_fov= 86, actor_to_attach= None, transform= None, # related to actor_to_attach ): camera_props = gymapi.CameraProperties() camera_props.enable_tensors = True camera_props.width = resolution[0] camera_props.height = resolution[1] camera_props.horizontal_fov = h_fov camera_handle = gym.create_camera_sensor(env_handle, camera_props) if actor_to_attach is not None: gym.attach_camera_to_body( camera_handle, env_handle, actor_to_attach, transform, gymapi.FOLLOW_POSITION, ) elif transform is not None: gym.set_camera_transform( camera_handle, env_handle, transform, ) return camera_handle @torch.no_grad() def play(args): env_cfg, train_cfg = task_registry.get_cfgs(name=args.task) with open(os.path.join("logs", train_cfg.runner.experiment_name, args.load_run, "config.json"), "r") as f: d = json.load(f, object_pairs_hook= OrderedDict) update_class_from_dict(env_cfg, d, strict= True) update_class_from_dict(train_cfg, d, strict= True) # override some parameters for testing if env_cfg.terrain.selected == "BarrierTrack": env_cfg.env.num_envs = min(env_cfg.env.num_envs, 1) env_cfg.env.episode_length_s = 5#20 env_cfg.terrain.max_init_terrain_level = 0 env_cfg.terrain.num_rows = 1 env_cfg.terrain.num_cols = 1 else: env_cfg.env.num_envs = min(env_cfg.env.num_envs, 1) env_cfg.env.episode_length_s = 60 env_cfg.terrain.terrain_length = 8 env_cfg.terrain.terrain_width = 8 env_cfg.terrain.max_init_terrain_level = 0 env_cfg.terrain.num_rows = 1 env_cfg.terrain.num_cols = 1 env_cfg.terrain.curriculum = False env_cfg.terrain.BarrierTrack_kwargs["options"] = [ # "crawl", # "jump", "leap", # "tilt", ] env_cfg.terrain.BarrierTrack_kwargs["leap"] = dict( length= (1.0, 1.0), depth= (0.4, 0.8), height= 0.2, ) if "one_obstacle_per_track" in env_cfg.terrain.BarrierTrack_kwargs.keys(): env_cfg.terrain.BarrierTrack_kwargs.pop("one_obstacle_per_track") env_cfg.terrain.BarrierTrack_kwargs["n_obstacles_per_track"] = 2# 2 env_cfg.commands.ranges.lin_vel_x = [1.5, 1.5] # [1.2, 1.2] env_cfg.terrain.BarrierTrack_kwargs['track_block_length']= 2.0 if "distill" in args.task: env_cfg.commands.ranges.lin_vel_x = [0.0, 0.0] env_cfg.commands.ranges.lin_vel_y = [-0., 0.] env_cfg.commands.ranges.ang_vel_yaw = [-0., 0.] env_cfg.domain_rand.push_robots = False env_cfg.domain_rand.init_base_pos_range = dict( x= [0.6, 0.6], y= [-0.05, 0.05], ) env_cfg.commands.curriculum = False env_cfg.termination.termination_terms = [] env_cfg.termination.timeout_at_border = False env_cfg.termination.timeout_at_finished = False env_cfg.viewer.debug_viz = False # in a1_distill, setting this to true will constantly showing the egocentric depth view. env_cfg.viewer.draw_volume_sample_points = False train_cfg.runner.resume = True train_cfg.runner_class_name = "OnPolicyRunner" if "distill" in args.task: # to save the memory train_cfg.algorithm.teacher_policy.sub_policy_paths = [] train_cfg.algorithm.teacher_policy_class_name = "ActorCritic" train_cfg.algorithm.teacher_policy = dict( num_actor_obs= 48, num_critic_obs= 48, num_actions= 12, ) ######### Some hacks to run ActorCriticMutex policy ########## if False: # for a1 train_cfg.runner.policy_class_name = "ActorCriticClimbMutex" train_cfg.policy.sub_policy_class_name = "ActorCriticRecurrent" logs_root = "logs" train_cfg.policy.sub_policy_paths = [ # must in the order of obstacle ID logs_root + "/field_a1_oracle/Jun03_00-01-38_SkillsPlaneWalking_pEnergySubsteps1e-5_rAlive2_pTorqueExceedIndicate1e+1_noCurriculum_propDelay0.04-0.05_noPerlinRate-2.0_nSubsteps4_envFreq50.0_aScale244", logs_root + "/field_a1_oracle/Aug08_05-22-52_Skills_tilt_pEnergySubsteps1e-5_rAlive1_pPenV5e-3_pPenD5e-3_pPosY0.50_pYaw0.50_rTilt7e-1_pTorqueExceedIndicate1e-1_virtualTerrain_propDelay0.04-0.05_push/", logs_root + "/field_a1_oracle/May21_05-25-19_Skills_crawl_pEnergy2e-5_rAlive1_pPenV6e-2_pPenD6e-2_pPosY0.2_kp50_noContactTerminate_aScale0.5/", logs_root + "/field_a1_oracle/Jun03_00-33-08_Skills_climb_pEnergySubsteps2e-6_rAlive2_pTorqueExceedIndicate2e-1_propDelay0.04-0.05_noPerlinRate0.2_nSubsteps4_envFreq50.0_aScale0.5", logs_root + "/field_a1_oracle/Jun04_01-03-59_Skills_leap_pEnergySubsteps2e-6_rAlive2_pPenV4e-3_pPenD4e-3_pPosY0.20_pYaw0.20_pTorqueExceedSquare1e-3_leapH0.2_propDelay0.04-0.05_noPerlinRate0.2_aScale0.5", ] train_cfg.policy.jump_down_policy_path = logs_root + "/field_a1_oracle/Aug30_16-12-14_Skills_climb_climbDownProb0.5_propDelay0.04-0.05" train_cfg.runner.resume = False env_cfg.env.use_lin_vel = True train_cfg.policy.cmd_vel_mapping = { 0: 1.0, 1: 0.5, 2: 0.8, 3: 1.2, 4: 1.5, } if args.task == "a1_distill": env_cfg.env.obs_components = env_cfg.env.privileged_obs_components env_cfg.env.privileged_obs_gets_privilege = False # prepare environment env, _ = task_registry.make_env(name=args.task, args=args, env_cfg=env_cfg) env.reset() print("terrain_levels:", env.terrain_levels.float().mean(), env.terrain_levels.float().max(), env.terrain_levels.float().min()) obs = env.get_observations() critic_obs = env.get_privileged_observations() # register debugging options to manually trigger disruption env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_P, "push_robot") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_L, "press_robot") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_J, "action_jitter") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_Q, "exit") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_R, "agent_full_reset") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_U, "full_reset") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_C, "resample_commands") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_W, "forward") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_S, "backward") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_A, "leftward") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_D, "rightward") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_F, "leftturn") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_G, "rightturn") env.gym.subscribe_viewer_keyboard_event(env.viewer, isaacgym.gymapi.KEY_X, "stop") # load policy ppo_runner, train_cfg = task_registry.make_alg_runner( env=env, name=args.task, args=args, train_cfg=train_cfg, save_cfg= False, ) agent_model = ppo_runner.alg.actor_critic policy = ppo_runner.get_inference_policy(device=env.device) ### get obs_slice to read the obs # obs_slice = get_obs_slice(env.obs_segments, "engaging_block") # 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') export_policy_as_jit(ppo_runner.alg.actor_critic, path) print('Exported policy as jit script to: ', path) if RECORD_FRAMES: transform = gymapi.Transform() transform.p = gymapi.Vec3(*env_cfg.viewer.pos) transform.r = gymapi.Quat.from_euler_zyx(0., 0., -np.pi/2) recording_camera = create_recording_camera( env.gym, env.envs[0], transform= transform, ) logger = Logger(env.dt) robot_index = 0 # which robot is used for logging joint_index = 4 # which joint is used for logging stop_state_log = 512 # number of steps before plotting states stop_rew_log = env.max_episode_length + 1 # number of steps before print average episode rewards camera_position = np.array(env_cfg.viewer.pos, dtype=np.float64) camera_vel = np.array([0.6, 0., 0.]) camera_direction = np.array(env_cfg.viewer.lookat) - np.array(env_cfg.viewer.pos) camera_follow_id = 0 # only effective when CAMERA_FOLLOW img_idx = 0 if hasattr(env, "motor_strength"): print("motor_strength:", env.motor_strength[robot_index].cpu().numpy().tolist()) print("torque_limits:", env.torque_limits) start_time = time.time_ns() for i in range(10*int(env.max_episode_length)): if "obs_slice" in locals().keys(): obs_component = obs[:, obs_slice[0]].reshape(-1, *obs_slice[1]) print(obs_component[robot_index]) actions = policy(obs.detach()) teacher_actions = actions obs, critic_obs, rews, dones, infos = env.step(actions.detach()) if RECORD_FRAMES: filename = os.path.join( os.path.abspath("logs/images/"), f"{img_idx:04d}.png", ) env.gym.write_viewer_image_to_file(env.viewer, filename) img_idx += 1 if MOVE_CAMERA: if CAMERA_FOLLOW: camera_position[:] = env.root_states[camera_follow_id, :3].cpu().numpy() - camera_direction else: camera_position += camera_vel * env.dt env.set_camera(camera_position, camera_position + camera_direction) for ui_event in env.gym.query_viewer_action_events(env.viewer): if ui_event.action == "push_robot" and ui_event.value > 0: # manully trigger to push the robot env._push_robots() if ui_event.action == "press_robot" and ui_event.value > 0: env.root_states[:, 9] = torch_rand_float(-env.cfg.domain_rand.max_push_vel_xy, 0, (env.num_envs, 1), device=env.device).squeeze(1) env.gym.set_actor_root_state_tensor(env.sim, gymtorch.unwrap_tensor(env.all_root_states)) if ui_event.action == "action_jitter" and ui_event.value > 0: # assuming wrong action is taken obs, critic_obs, rews, dones, infos = env.step(torch.tanh(torch.randn_like(actions))) if ui_event.action == "exit" and ui_event.value > 0: print("exit") exit(0) if ui_event.action == "agent_full_reset" and ui_event.value > 0: print("agent_full_reset") agent_model.reset() if ui_event.action == "full_reset" and ui_event.value > 0: print("full_reset") agent_model.reset() obs, _ = env.reset() if ui_event.action == "resample_commands" and ui_event.value > 0: print("resample_commands") env._resample_commands(torch.arange(env.num_envs, device= env.device)) if ui_event.action == "stop" and ui_event.value > 0: env.commands[:, :] = 0 if ui_event.action == "forward" and ui_event.value > 0: env.commands[:, 0] = env_cfg.commands.ranges.lin_vel_x[1] if ui_event.action == "backward" and ui_event.value > 0: env.commands[:, 0] = env_cfg.commands.ranges.lin_vel_x[0] if ui_event.action == "leftward" and ui_event.value > 0: env.commands[:, 1] = env_cfg.commands.ranges.lin_vel_y[1] if ui_event.action == "rightward" and ui_event.value > 0: env.commands[:, 1] = env_cfg.commands.ranges.lin_vel_y[0] if ui_event.action == "leftturn" and ui_event.value > 0: env.commands[:, 2] = env_cfg.commands.ranges.ang_vel_yaw[1] if ui_event.action == "rightturn" and ui_event.value > 0: env.commands[:, 2] = env_cfg.commands.ranges.ang_vel_yaw[0] if i < stop_state_log: if torch.is_tensor(env.cfg.control.action_scale): action_scale = env.cfg.control.action_scale.detach().cpu().numpy()[joint_index] else: action_scale = env.cfg.control.action_scale base_roll = get_euler_xyz(env.base_quat)[0][robot_index].item() base_pitch = get_euler_xyz(env.base_quat)[1][robot_index].item() if base_pitch > torch.pi: base_pitch -= torch.pi * 2 logger.log_states( { 'dof_pos_target': actions[robot_index, joint_index].item() * action_scale, 'dof_pos': (env.dof_pos - env.default_dof_pos)[robot_index, joint_index].item(), 'dof_vel': env.substep_dof_vel[robot_index, 0, joint_index].max().item(), 'dof_torque': env.substep_torques[robot_index, 0, joint_index].max().item(), 'command_x': env.commands[robot_index, 0].item(), 'command_y': env.commands[robot_index, 1].item(), 'command_yaw': env.commands[robot_index, 2].item(), 'base_vel_x': env.base_lin_vel[robot_index, 0].item(), 'base_vel_y': env.base_lin_vel[robot_index, 1].item(), 'base_vel_z': env.base_lin_vel[robot_index, 2].item(), 'base_pitch': base_pitch, 'contact_forces_z': env.contact_forces[robot_index, env.feet_indices, 2].cpu().numpy(), 'max_torques': torch.abs(env.substep_torques).max().item(), "student_action": actions[robot_index, 2].item(), "teacher_action": teacher_actions[robot_index, 2].item(), "reward": rews[robot_index].item(), } ) elif i==stop_state_log: logger.plot_states() env._get_terrain_curriculum_move(torch.tensor([0], device= env.device)) if 0 < i < stop_rew_log: if infos["episode"]: num_episodes = torch.sum(env.reset_buf).item() if num_episodes>0: logger.log_rewards(infos["episode"], num_episodes) elif i==stop_rew_log: logger.print_rewards() if dones.any(): agent_model.reset(dones) print("env dones,{} because has timeout".format("" if env.time_out_buf[dones].any() else " not")) print(infos) if i % 100 == 0: print("frame_rate:" , 100/(time.time_ns() - start_time) * 1e9, "command_x:", env.commands[robot_index, 0], ) start_time = time.time_ns() if __name__ == '__main__': EXPORT_POLICY = False RECORD_FRAMES = False MOVE_CAMERA = True CAMERA_FOLLOW = True args = get_args() play(args)