diff --git a/README.md b/README.md index cfd81a9..ca5f744 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,7 @@ This repository provides a set of reinforcement learning tasks for Booster robots using [Isaac Lab](https://isaac-sim.github.io/IsaacLab/main/index.html). Currently it includes the fabulous [BeyondMimic motion tracking](https://github.com/HybridRobotics/whole_body_tracking) framework adapted to Booster K1 robots. +The motion conversion and replay utilities under `scripts/mimic/` support both Booster K1 and T1 robots. This repository follows the standard Isaac Lab project structure, and is tested with IsaacLab 2.2 and Isaac Sim 5.0. ## Installation @@ -30,9 +31,15 @@ This repository follows the standard Isaac Lab project structure, and is tested - Prepare BeyondMimic motion data: ```bash # use 'FULL_PATH_TO_isaaclab.sh|bat -p' instead of 'python' if Isaac Lab is not installed in Python venv or conda - python scripts/csv_to_npz.py --headless --input_file=/motions/K1/.csv --input_fps= --output_name=/motions/K1/.npz + python scripts/mimic/csv_to_npz.py --headless --input_file=/motions//.csv --input_fps= --output_file=/motions//.npz --output_fps=50 --robot= ``` + Optional arguments: + + - `--robot=t1` converts T1 motion files. The default robot is `k1`. + - `--frame_range ` converts only a subset of frames. Frame indices are 1-based and inclusive. + - `--output_fps` controls the interpolation rate of the exported `.npz` motion. + ## Usage - Listing the available tasks: @@ -58,6 +65,17 @@ This repository follows the standard Isaac Lab project structure, and is tested This script also exports the trained policy to a TorchScript/ONNX file for deployment on real robots in `logs/rsl_rl///exported/`. +- Replay a converted motion file for inspection: + + ```bash + python scripts/mimic/replay_npz.py --motion=.npz --robot= + + # or download and replay from Weights & Biases registry + python scripts/mimic/replay_npz.py --registry_name= --robot= + ``` + + When `--registry_name` does not include an alias, the script automatically uses `:latest`. + ## Deploy After a model has been trained and exported, you can deploy the trained policy in MuJoCo or on real Booster robots using the [booster_deploy](https://github.com/BoosterRobotics/booster_deploy) repository. For more details, please refer to the instructions in the [booster_deploy](https://github.com/BoosterRobotics/booster_deploy) repository. diff --git a/scripts/csv_to_npz.py b/scripts/mimic/csv_to_npz.py similarity index 91% rename from scripts/csv_to_npz.py rename to scripts/mimic/csv_to_npz.py index c1a2405..beccfe9 100644 --- a/scripts/csv_to_npz.py +++ b/scripts/mimic/csv_to_npz.py @@ -4,7 +4,7 @@ # Usage python csv_to_npz.py --input_file LAFAN/dance1_subject2.csv --input_fps 30 --frame_range 122 722 \ - --output_file ./motions/dance1_subject2.npz --output_fps 50 + --output_file ./motions/dance1_subject2.npz --output_fps 50 --robot k1 """ """Launch Isaac Sim Simulator first.""" @@ -29,8 +29,15 @@ " loaded." ), ) -parser.add_argument("--output_name", type=str, required=True, help="The name of the motion npz file.") +parser.add_argument("--output_file", type=str, required=True, help="The name of the motion npz file.") parser.add_argument("--output_fps", type=int, default=50, help="The fps of the output motion.") +parser.add_argument( + "--robot", + type=str.lower, + default="k1", + choices=["k1", "t1"], + help="Target robot type. Defaults to k1. Use t1 to convert T1 motion format.", +) # append AppLauncher cli args AppLauncher.add_app_launcher_args(parser) @@ -57,8 +64,14 @@ # Pre-defined configs ## # Booster assets currently only provide retargeted motion data for Booster K1 robot -from booster_train.assets.robots.booster import BOOSTER_K1_CFG as ROBOT_CFG -from booster_assets.motions import K1_JOINT_NAMES as JOINT_NAMES +from booster_train.assets.robots.booster import BOOSTER_K1_CFG, BOOSTER_T1_CFG +from booster_assets.motions import K1_JOINT_NAMES, T1_JOINT_NAMES + + +ROBOT_CONFIGS = { + "k1": (BOOSTER_K1_CFG, K1_JOINT_NAMES), + "t1": (BOOSTER_T1_CFG, T1_JOINT_NAMES), +} @configclass @@ -78,7 +91,7 @@ class ReplayMotionsSceneCfg(InteractiveSceneCfg): ) # articulation - robot: ArticulationCfg = ROBOT_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + robot: ArticulationCfg = BOOSTER_K1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") class MotionLoader: @@ -194,12 +207,15 @@ def _so3_derivative(self, rotations: torch.Tensor, dt: float) -> torch.Tensor: def get_next_state( self, ) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ], + bool, ]: """Gets the next state of the motion.""" state = ( @@ -303,19 +319,22 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene, joi ): log[k] = np.stack(log[k], axis=0) - np.savez(f"{args_cli.output_name}", **log) - print(f"[INFO]: Motion saved to {args_cli.output_name}") + np.savez(f"{args_cli.output_file}", **log) + print(f"[INFO]: Motion saved to {args_cli.output_file}") sys.exit(0) def main(): """Main function.""" + robot_cfg, joint_names = ROBOT_CONFIGS[args_cli.robot] + # Load kit helper sim_cfg = sim_utils.SimulationCfg(device=args_cli.device) sim_cfg.dt = 1.0 / args_cli.output_fps sim = SimulationContext(sim_cfg) # Design scene scene_cfg = ReplayMotionsSceneCfg(num_envs=1, env_spacing=2.0) + scene_cfg.robot = robot_cfg.replace(prim_path="{ENV_REGEX_NS}/Robot") scene = InteractiveScene(scene_cfg) # Play the simulator sim.reset() @@ -325,7 +344,7 @@ def main(): run_simulator( sim, scene, - joint_names=JOINT_NAMES, + joint_names=joint_names, ) diff --git a/scripts/replay_npz.py b/scripts/mimic/replay_npz.py similarity index 79% rename from scripts/replay_npz.py rename to scripts/mimic/replay_npz.py index 6d62ae2..2dc738a 100644 --- a/scripts/replay_npz.py +++ b/scripts/mimic/replay_npz.py @@ -1,12 +1,12 @@ -"""This script demonstrates how to replay K1 robot motions from npz files. +"""This script demonstrates how to replay robot motions from npz files. .. code-block:: bash # Usage - Direct file path - python replay_npz.py --motion + python replay_npz.py --motion --robot k1 # Usage - From wandb registry - python replay_npz.py --registry_name + python replay_npz.py --registry_name --robot t1 """ """Launch Isaac Sim Simulator first.""" @@ -23,6 +23,13 @@ parser = argparse.ArgumentParser(description="Replay converted motions.") parser.add_argument("--motion", type=str, default=None, help="Path to the motion npz file.") parser.add_argument("--registry_name", type=str, default=None, help="The name of the wand registry.") +parser.add_argument( + "--robot", + type=str.lower, + default="k1", + choices=["k1", "t1"], + help="Target robot type. Defaults to k1. Use t1 to replay T1 motion format.", +) # append AppLauncher cli args AppLauncher.add_app_launcher_args(parser) @@ -45,10 +52,16 @@ ## # Pre-defined configs ## -from booster_train.assets.robots.booster import BOOSTER_K1_CFG +from booster_train.assets.robots.booster import BOOSTER_K1_CFG, BOOSTER_T1_CFG from booster_train.tasks.manager_based.beyond_mimic.mdp.commands import MotionLoader +ROBOT_CONFIGS = { + "k1": BOOSTER_K1_CFG, + "t1": BOOSTER_T1_CFG, +} + + @configclass class ReplayMotionsSceneCfg(InteractiveSceneCfg): """Configuration for a replay motions scene.""" @@ -95,15 +108,18 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): else: raise ValueError("Either --motion or --registry_name must be provided.") - # Load npz file to get body names and determine body_indexes - # For K1, we typically use Trunk as anchor body (index 0) - # body_indexes should be a list of indices corresponding to the bodies we want to use - # For replay, we only need the anchor body (Trunk), which is typically at index 0 - body_indexes = [0] # Default to index 0 for anchor body (Trunk) + with np.load(motion_file) as motion_data: + motion_body_names = motion_data["body_names"].tolist() if "body_names" in motion_data else robot.body_names + + # For both K1/T1 converted motions, the anchor body is stored at index 0. + anchor_body_name = motion_body_names[0] motion = MotionLoader( motion_file, - body_indexes, + [anchor_body_name], + robot.joint_names, + default_motion_body_names=robot.body_names, + default_motion_joint_names=robot.joint_names, tail_len=0, device=str(sim.device), ) @@ -116,7 +132,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): time_steps[reset_ids] = 0 root_states = robot.data.default_root_state.clone() - root_states[:, :3] = motion.body_pos_w[time_steps][:, 0] + scene.env_origins[:, None, :] + root_states[:, :3] = motion.body_pos_w[time_steps][:, 0] + root_states[:, :2] += scene.env_origins[:, :2] root_states[:, 3:7] = motion.body_quat_w[time_steps][:, 0] root_states[:, 7:10] = motion.body_lin_vel_w[time_steps][:, 0] root_states[:, 10:] = motion.body_ang_vel_w[time_steps][:, 0] @@ -132,11 +149,14 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): def main(): + robot_cfg = ROBOT_CONFIGS[args_cli.robot] + sim_cfg = sim_utils.SimulationCfg(device=args_cli.device) sim_cfg.dt = 0.02 sim = SimulationContext(sim_cfg) scene_cfg = ReplayMotionsSceneCfg(num_envs=1, env_spacing=2.0) + scene_cfg.robot = robot_cfg.replace(prim_path="{ENV_REGEX_NS}/Robot") scene = InteractiveScene(scene_cfg) sim.reset() # Run the simulator diff --git a/source/booster_train/booster_train/assets/robots/booster.py b/source/booster_train/booster_train/assets/robots/booster.py index 47ced15..c366388 100644 --- a/source/booster_train/booster_train/assets/robots/booster.py +++ b/source/booster_train/booster_train/assets/robots/booster.py @@ -329,3 +329,19 @@ ), }, ) + +T1_ACTION_SCALE = {} +for a in BOOSTER_T1_CFG.actuators.values(): + e = a.effort_limit_sim + s = a.stiffness + names = a.joint_names_expr + if not isinstance(e, dict): + e = {n: e for n in names} + if not isinstance(s, dict): + s = {n: s for n in names} + for n in names: + if n in e and n in s and s[n]: + T1_ACTION_SCALE[n] = 0.25 * e[n] / s[n] + +print(f'{BOOSTER_T1_CFG.actuators=}') +print(f'{T1_ACTION_SCALE=}') \ No newline at end of file diff --git a/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/__init__.py b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/__init__.py b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/__init__.py new file mode 100644 index 0000000..eeb93b6 --- /dev/null +++ b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/__init__.py @@ -0,0 +1,30 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import gymnasium as gym + +## +# Register Gym environments. +## + +gym.register( + id="Booster-T1-Dance-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.env_cfg:RoughWoStateEstimationEnvCfg", + "rsl_rl_cfg_entry_point": f"{__name__}.ppo_cfg:PPORunnerCfg", + }, +) + +gym.register( + id="Booster-T1-Dance-v0-Play", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.env_cfg:PlayFlatWoStateEstimationEnvCfg", + "rsl_rl_cfg_entry_point": f"{__name__}.ppo_cfg:PPORunnerCfg", + }, +) \ No newline at end of file diff --git a/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/env_cfg.py b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/env_cfg.py new file mode 100644 index 0000000..ee5b884 --- /dev/null +++ b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/env_cfg.py @@ -0,0 +1,87 @@ +from isaaclab.utils import configclass +from isaaclab.terrains import TerrainGeneratorCfg +import isaaclab.terrains as terrain_gen +from booster_assets import BOOSTER_ASSETS_DIR +from booster_train.assets.robots.booster import BOOSTER_T1_CFG as ROBOT_CFG, T1_ACTION_SCALE +from booster_train.tasks.manager_based.beyond_mimic.agents.rsl_rl_ppo_cfg import LOW_FREQ_SCALE +from .tracking_env_cfg import TrackingEnvCfg + + +@configclass +class FlatEnvCfg(TrackingEnvCfg): + def __post_init__(self): + super().__post_init__() + + self.scene.robot = ROBOT_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + self.actions.joint_pos.scale = T1_ACTION_SCALE + self.commands.motion.motion_file = f"{BOOSTER_ASSETS_DIR}/motions/T1/CLIO_Dimitroula_1_stageii.npz" + self.commands.motion.anchor_body_name = "Trunk" + self.commands.motion.body_names = [ + 'Trunk', + 'H2', + 'Hip_Roll_Left', + 'Shank_Left', + 'left_foot_link', + 'Hip_Roll_Right', + 'Shank_Right', + 'right_foot_link', + 'Waist', + 'AL2', + 'AL3', + 'left_hand_link', + 'AR2', + 'AR3', + 'right_hand_link', + ] + + +@configclass +class FlatWoStateEstimationEnvCfg(FlatEnvCfg): + def __post_init__(self): + super().__post_init__() + self.observations.policy.motion_anchor_pos_b = None + self.observations.policy.base_lin_vel = None + + +@configclass +class RoughWoStateEstimationEnvCfg(FlatWoStateEstimationEnvCfg): + def __post_init__(self): + super().__post_init__() + + self.scene.terrain.terrain_type = "generator" + self.scene.terrain.debug_vis = False # 设为True可视化地形分布 + self.scene.terrain.terrain_generator = TerrainGeneratorCfg( + size=(10.0, 10.0), # 每个地形块尺寸（米） + border_width=20.0, # 边界宽度（米） + num_rows=5, # 地形网格行数 + num_cols=10, # 地形网格列数 + horizontal_scale=0.1, # 水平分辨率 + vertical_scale=0.005, # 垂直分辨率 + slope_threshold=0.75, # 网格简化阈值 + use_cache=False, # 每次重新生成地形 + curriculum=False, # 启用课程学习 + sub_terrains={ + # 80%接近平面的地形（非常平滑） + "nearly_flat": terrain_gen.HfRandomUniformTerrainCfg( + proportion=0.8, + noise_range=(0.0, 0.005), # 高度波动0-0.5cm（几乎平坦） + noise_step=0.005, # 噪声步长0.5cm + border_width=0.25, + ), + # 20%随机粗糙地形 + "random_rough": terrain_gen.HfRandomUniformTerrainCfg( + proportion=0.2, + noise_range=(-0.015, 0.015), # 高度波动±1.5cm + noise_step=0.005, # 噪声步长0.5cm + border_width=0.25, + ), + }, + ) + + +@configclass +class PlayFlatWoStateEstimationEnvCfg(FlatWoStateEstimationEnvCfg): + def __post_init__(self): + super().__post_init__() + self.commands.motion.play = True + self.events.push_robot = None diff --git a/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/ppo_cfg.py b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/ppo_cfg.py new file mode 100644 index 0000000..7d82b3e --- /dev/null +++ b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/ppo_cfg.py @@ -0,0 +1,8 @@ +from isaaclab.utils import configclass +from booster_train.tasks.manager_based.beyond_mimic.agents.rsl_rl_ppo_cfg import BasePPORunnerCfg + + +@configclass +class PPORunnerCfg(BasePPORunnerCfg): + max_iterations = 30000 + experiment_name = "t1_dance" diff --git a/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/tracking_env_cfg.py b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/tracking_env_cfg.py new file mode 100644 index 0000000..2c0d4f9 --- /dev/null +++ b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/T1/dance/tracking_env_cfg.py @@ -0,0 +1,323 @@ +from __future__ import annotations + +from dataclasses import MISSING + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, AssetBaseCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sensors import ContactSensorCfg +from isaaclab.terrains import TerrainImporterCfg + +## +# Pre-defined configs +## +from isaaclab.utils import configclass +from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise + +import booster_train.tasks.manager_based.beyond_mimic.mdp as mdp + +## +# Scene definition +## + +VELOCITY_RANGE = { + "x": (-0.5, 0.5), + "y": (-0.5, 0.5), + "z": (-0.2, 0.2), + "roll": (-0.52, 0.52), + "pitch": (-0.52, 0.52), + "yaw": (-0.78, 0.78), +} + + +@configclass +class MySceneCfg(InteractiveSceneCfg): + """Configuration for the terrain scene with a legged robot.""" + + # ground terrain + terrain = TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="plane", + collision_group=-1, + physics_material=sim_utils.RigidBodyMaterialCfg( + friction_combine_mode="multiply", + restitution_combine_mode="multiply", + static_friction=1.0, + dynamic_friction=1.0, + ), + visual_material=sim_utils.MdlFileCfg( + mdl_path="{NVIDIA_NUCLEUS_DIR}/Materials/Base/Architecture/Shingles_01.mdl", + project_uvw=True, + ), + ) + # robots + robot: ArticulationCfg = MISSING + # lights + light = AssetBaseCfg( + prim_path="/World/light", + spawn=sim_utils.DistantLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0), + ) + sky_light = AssetBaseCfg( + prim_path="/World/skyLight", + spawn=sim_utils.DomeLightCfg(color=(0.13, 0.13, 0.13), intensity=1000.0), + ) + contact_forces = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, track_air_time=True, force_threshold=10.0, debug_vis=True + ) + + +## +# MDP settings +## + + +@configclass +class CommandsCfg: + """Command specifications for the MDP.""" + + motion = mdp.MotionCommandCfg( + asset_name="robot", + resampling_time_range=(1.0e9, 1.0e9), + debug_vis=True, + pose_range={ + "x": (-0.05, 0.05), + "y": (-0.05, 0.05), + "z": (-0.01, 0.01), + "roll": (-0.1, 0.1), + "pitch": (-0.1, 0.1), + "yaw": (-0.2, 0.2), + }, + velocity_range=VELOCITY_RANGE, + joint_position_range=(-0.1, 0.1), + ) + + +@configclass +class ActionsCfg: + """Action specifications for the MDP.""" + + joint_pos = mdp.JointPositionActionCfg(asset_name="robot", joint_names=[".*"], use_default_offset=True) + + +@configclass +class ObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group.""" + + # observation terms (order preserved) + command = ObsTerm(func=mdp.generated_commands, params={"command_name": "motion"}) + motion_anchor_pos_b = ObsTerm( + func=mdp.motion_anchor_pos_b, params={"command_name": "motion"}, noise=Unoise(n_min=-0.25, n_max=0.25) + ) + motion_anchor_ori_b = ObsTerm( + func=mdp.motion_anchor_ori_b, params={"command_name": "motion"}, noise=Unoise(n_min=-0.05, n_max=0.05) + ) + base_lin_vel = ObsTerm(func=mdp.base_lin_vel, noise=Unoise(n_min=-0.5, n_max=0.5)) + base_ang_vel = ObsTerm(func=mdp.base_ang_vel, noise=Unoise(n_min=-0.2, n_max=0.2)) + joint_pos = ObsTerm(func=mdp.joint_pos_rel, noise=Unoise(n_min=-0.01, n_max=0.01)) + joint_vel = ObsTerm(func=mdp.joint_vel_rel, noise=Unoise(n_min=-0.5, n_max=0.5)) + actions = ObsTerm(func=mdp.last_action) + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_terms = True + + @configclass + class PrivilegedCfg(ObsGroup): + command = ObsTerm(func=mdp.generated_commands, params={"command_name": "motion"}) + motion_anchor_pos_b = ObsTerm(func=mdp.motion_anchor_pos_b, params={"command_name": "motion"}) + motion_anchor_ori_b = ObsTerm(func=mdp.motion_anchor_ori_b, params={"command_name": "motion"}) + body_pos = ObsTerm(func=mdp.robot_body_pos_b, params={"command_name": "motion"}) + body_ori = ObsTerm(func=mdp.robot_body_ori_b, params={"command_name": "motion"}) + base_lin_vel = ObsTerm(func=mdp.base_lin_vel) + base_ang_vel = ObsTerm(func=mdp.base_ang_vel) + joint_pos = ObsTerm(func=mdp.joint_pos_rel) + joint_vel = ObsTerm(func=mdp.joint_vel_rel) + actions = ObsTerm(func=mdp.last_action) + + # observation groups + policy: PolicyCfg = PolicyCfg() + critic: PrivilegedCfg = PrivilegedCfg() + + +@configclass +class EventCfg: + """Configuration for events.""" + + # startup + physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*"), + "static_friction_range": (0.3, 0.6), + "dynamic_friction_range": (0.3, 0.6), + "restitution_range": (0.0, 0.5), + "num_buckets": 64, + }, + ) + + add_joint_default_pos = EventTerm( + func=mdp.randomize_joint_default_pos, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=[".*"]), + "pos_distribution_params": (-0.01, 0.01), + "operation": "add", + }, + ) + + base_com = EventTerm( + func=mdp.randomize_rigid_body_com, + mode="startup", + params={ + # "asset_cfg": SceneEntityCfg("robot", body_names="torso_link"), + "asset_cfg": SceneEntityCfg("robot", body_names="Trunk"), + "com_range": {"x": (-0.025, 0.025), "y": (-0.05, 0.05), "z": (-0.05, 0.05)}, + }, + ) + + # interval + push_robot = EventTerm( + func=mdp.push_by_setting_velocity, + mode="interval", + interval_range_s=(1.0, 3.0), + params={"velocity_range": VELOCITY_RANGE}, + ) + + +@configclass +class RewardsCfg: + """Reward terms for the MDP.""" + + motion_global_anchor_pos = RewTerm( + func=mdp.motion_global_anchor_position_error_exp, + weight=0.5, + params={"command_name": "motion", "std": 0.3}, + ) + motion_global_anchor_ori = RewTerm( + func=mdp.motion_global_anchor_orientation_error_exp, + weight=0.5, + params={"command_name": "motion", "std": 0.4}, + ) + motion_body_pos = RewTerm( + func=mdp.motion_relative_body_position_error_exp, + weight=1.0, + params={"command_name": "motion", "std": 0.3}, + ) + motion_body_ori = RewTerm( + func=mdp.motion_relative_body_orientation_error_exp, + weight=1.0, + params={"command_name": "motion", "std": 0.4}, + ) + motion_body_lin_vel = RewTerm( + func=mdp.motion_global_body_linear_velocity_error_exp, + weight=1.0, + params={"command_name": "motion", "std": 1.0}, + ) + motion_body_ang_vel = RewTerm( + func=mdp.motion_global_body_angular_velocity_error_exp, + weight=1.0, + params={"command_name": "motion", "std": 3.14}, + ) + action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-1e-1) + joint_limit = RewTerm( + func=mdp.joint_pos_limits, + weight=-10.0, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])}, + ) + undesired_contacts = RewTerm( + func=mdp.undesired_contacts, + weight=-0.1, + params={ + "sensor_cfg": SceneEntityCfg( + "contact_forces", + body_names=[ + r"^(?!left_hand_link$)(?!right_hand_link$)(?!left_foot_link$)(?!right_foot_link$).+$" + ], + ), + "threshold": 1.0, + }, + ) + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out = DoneTerm(func=mdp.time_out, time_out=True) + anchor_pos = DoneTerm( + func=mdp.bad_anchor_pos_z_only, + params={"command_name": "motion", "threshold": 0.25}, + ) + anchor_ori = DoneTerm( + func=mdp.bad_anchor_ori, + params={"asset_cfg": SceneEntityCfg("robot"), "command_name": "motion", "threshold": 0.8}, + ) + ee_body_pos = DoneTerm( + func=mdp.bad_motion_body_pos_z_only, + params={ + "command_name": "motion", + "threshold": 0.25, + "body_names": [ + "left_hand_link", + "right_hand_link", + "left_foot_link", + "right_foot_link", + ], + }, + ) + + +@configclass +class CurriculumCfg: + """Curriculum terms for the MDP.""" + + pass + + +## +# Environment configuration +## + + +@configclass +class TrackingEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the locomotion velocity-tracking environment.""" + + # Scene settings + scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=2.5) + # Basic settings + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + # MDP settings + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + curriculum: CurriculumCfg = CurriculumCfg() + + def __post_init__(self): + """Post initialization.""" + # general settings + self.decimation = 4 + self.episode_length_s = 10.0 + # simulation settings + self.sim.dt = 0.005 + self.sim.render_interval = self.decimation + self.sim.physics_material = self.scene.terrain.physics_material + self.sim.physx.gpu_max_rigid_patch_count = 10 * 2**15 + # viewer settings + self.viewer.eye = (1.5, 1.5, 1.5) + self.viewer.origin_type = "asset_root" + self.viewer.asset_name = "robot" diff --git a/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/k1/mj_dance_004/__init__.py b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/k1/mj_dance_004/__init__.py index 2f87207..410dc32 100644 --- a/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/k1/mj_dance_004/__init__.py +++ b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/k1/mj_dance_004/__init__.py @@ -18,3 +18,13 @@ "rsl_rl_cfg_entry_point": f"{__name__}.ppo_cfg:PPORunnerCfg", }, ) + +gym.register( + id="Booster-K1-MJ_Dance_004-v0-Play", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.env_cfg:PlayFlatWoStateEstimationEnvCfg", + "rsl_rl_cfg_entry_point": f"{__name__}.ppo_cfg:PPORunnerCfg", + }, +) \ No newline at end of file diff --git a/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/k1/mj_dance_004/env_cfg.py b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/k1/mj_dance_004/env_cfg.py index cb6fd94..7aabc6a 100644 --- a/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/k1/mj_dance_004/env_cfg.py +++ b/source/booster_train/booster_train/tasks/manager_based/beyond_mimic/robots/k1/mj_dance_004/env_cfg.py @@ -83,4 +83,12 @@ class FlatLowFreqEnvCfg(FlatEnvCfg): def __post_init__(self): super().__post_init__() self.decimation = round(self.decimation / LOW_FREQ_SCALE) - self.rewards.action_rate_l2.weight *= LOW_FREQ_SCALE \ No newline at end of file + self.rewards.action_rate_l2.weight *= LOW_FREQ_SCALE + + +@configclass +class PlayFlatWoStateEstimationEnvCfg(FlatWoStateEstimationEnvCfg): + def __post_init__(self): + super().__post_init__() + self.commands.motion.play = True + self.events.push_robot = None