diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 79baea5a5..d69308bb3 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -13,7 +13,7 @@ repos: - id: check-yaml - id: check-toml - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.14.5 + rev: v0.17.2 hooks: - id: ruff args: [ --fix ] diff --git a/docs/source/FAQ/index.md b/docs/source/FAQ/index.md index bdfaccacb..f888ff9fb 100644 --- a/docs/source/FAQ/index.md +++ b/docs/source/FAQ/index.md @@ -122,8 +122,8 @@ from metasim.sim.sapien_handler import SapienHandler from metasim.sim.isaacsim_handler import IsaacSimHandler # Same scenario works with any simulator -handler = MujocoHandler(scenario) # or -handler = SapienHandler(scenario) # or +handler = MujocoHandler(scenario) # or +handler = SapienHandler(scenario) # or handler = IsaacSimHandler(scenario) ``` diff --git a/docs/source/dataset_benchmark/benchmark/usage.md b/docs/source/dataset_benchmark/benchmark/usage.md index 12acacf3c..1a43e2478 100644 --- a/docs/source/dataset_benchmark/benchmark/usage.md +++ b/docs/source/dataset_benchmark/benchmark/usage.md @@ -43,16 +43,8 @@ class RandomizationCfg: Without the need for knowing the details of randomized evaluation, you can simply define the level of randomization with numbers 0-3, and use the randomization config inside of a scenario config: ```python -randomization = RandomizationCfg( - camera=False, light=False, ground=False, reflection=False -) -scenario = ScenarioCfg( - task=task, - robot=robot, - cameras=[camera], - randomization=randomization, - try_add_table=True -) +randomization = RandomizationCfg(camera=False, light=False, ground=False, reflection=False) +scenario = ScenarioCfg(task=task, robot=robot, cameras=[camera], randomization=randomization, try_add_table=True) ``` And the instantiated environment will be automatically randomized! diff --git a/docs/source/dataset_benchmark/dataset/multiagent.md b/docs/source/dataset_benchmark/dataset/multiagent.md index e3caa449a..6b9ba95bb 100644 --- a/docs/source/dataset_benchmark/dataset/multiagent.md +++ b/docs/source/dataset_benchmark/dataset/multiagent.md @@ -16,7 +16,7 @@ demos; each demo carries `init_state`, `actions`, and optional `states`: ```python { - "franka_left": [{"init_state": {...}, "actions": [...], "states": None}, ...], + "franka_left": [{"init_state": {...}, "actions": [...], "states": None}, ...], "franka_right": [{"init_state": {...}, "actions": [...], "states": None}, ...], "metadata": {"num_agents": 2, "agents": ["franka_left", "franka_right"]}, } diff --git a/docs/source/dataset_benchmark/integrations/libero.md b/docs/source/dataset_benchmark/integrations/libero.md index 8e4d29080..7a302730e 100644 --- a/docs/source/dataset_benchmark/integrations/libero.md +++ b/docs/source/dataset_benchmark/integrations/libero.md @@ -70,13 +70,13 @@ user-set value or the base `~/.libero`). No manual config editing needed. ## Usage ```python -import roboverse_pack.tasks.libero # auto-registers Libero/__ -import roboverse_pack.tasks.libero_plus # auto-registers LiberoPlus/__ +import roboverse_pack.tasks.libero # auto-registers Libero/__ +import roboverse_pack.tasks.libero_plus # auto-registers LiberoPlus/__ from roboverse_pack.tasks.libero_plus import make_liberoplus_env # build any of the 10,120 perturbation tasks by (suite, task index) env = make_liberoplus_env("libero_object", 7, seed=0) -obs, reward, done, info = env.step([0.0] * 7) # native legacy-gym 4-tuple (kept for fidelity) +obs, reward, done, info = env.step([0.0] * 7) # native legacy-gym 4-tuple (kept for fidelity) ``` ## Native MetaSim tasks — run (and delete) LIBERO @@ -93,9 +93,9 @@ other RoboVerse task and run with `libero`/`robosuite` uninstalled: from metasim.task.registry import get_task_class Task = get_task_class("libero_native.10__KITCHEN_SCENE3_turn_on_the_stove_and_put_the_moka_pot_on_it") -env = Task() # loads the vendored scene; no libero +env = Task() # loads the vendored scene; no libero obs, info = env.reset() -obs, reward, terminated, time_out, info = env.step(action) # 7-D OSC delta + gripper +obs, reward, terminated, time_out, info = env.step(action) # 7-D OSC delta + gripper ``` Each task is a first-class `BaseTaskEnv` (`NativeLiberoEnv`); a static bundle under diff --git a/docs/source/dataset_benchmark/integrations/mjlab.md b/docs/source/dataset_benchmark/integrations/mjlab.md index ea5527f65..c22a54f63 100644 --- a/docs/source/dataset_benchmark/integrations/mjlab.md +++ b/docs/source/dataset_benchmark/integrations/mjlab.md @@ -84,13 +84,15 @@ paths: ```python ScenarioCfg( - robots=[RobotCfg( - name="g1", - num_joints=29, - fix_base_link=False, - mjcf_path=mjlab_asset("asset_zoo/robots/unitree_g1/xmls/g1.xml"), - enabled_self_collisions=False, - )], + robots=[ + RobotCfg( + name="g1", + num_joints=29, + fix_base_link=False, + mjcf_path=mjlab_asset("asset_zoo/robots/unitree_g1/xmls/g1.xml"), + enabled_self_collisions=False, + ) + ], simulator="mujoco", add_default_ground=False, # raw rollout doesn't add a ground either ) diff --git a/docs/source/dataset_benchmark/integrations/simpler_env.md b/docs/source/dataset_benchmark/integrations/simpler_env.md index 9f911c76b..589598beb 100644 --- a/docs/source/dataset_benchmark/integrations/simpler_env.md +++ b/docs/source/dataset_benchmark/integrations/simpler_env.md @@ -69,20 +69,23 @@ rendering works, no sm_120 wall. ## Usage ```python -import roboverse_pack.tasks.simpler_env # auto-registers SimplerEnv/ + simpler. +import roboverse_pack.tasks.simpler_env # auto-registers SimplerEnv/ + simpler. # (1) MetaSim-native via gym import gymnasium as gym + env = gym.make("SimplerEnv/google_robot_pick_coke_can") obs, info = env.reset(seed=0) obs, reward, terminated, truncated, info = env.step(env.action_space.sample()) # (2) MetaSim-native via the MetaSim task registry from metasim.task.registry import get_task_class + task = get_task_class("simpler.widowx_stack_cube")() # (3) optional upstream passthrough (requires the SimplerEnv clone) from roboverse_pack.tasks.simpler_env import register_simpler_env_passthrough + register_simpler_env_passthrough(prefix="SimplerEnvPassthrough/") env = gym.make("SimplerEnvPassthrough/google_robot_pick_coke_can") ``` diff --git a/docs/source/roboverse_learn/imitation_learning/contributing.md b/docs/source/roboverse_learn/imitation_learning/contributing.md index 36609ca8d..8260eb0b9 100644 --- a/docs/source/roboverse_learn/imitation_learning/contributing.md +++ b/docs/source/roboverse_learn/imitation_learning/contributing.md @@ -136,6 +136,7 @@ from roboverse_learn.base.base_runner import BaseRunner abs_config_path = str(pathlib.Path(__file__).resolve().parent.joinpath("configs").absolute()) OmegaConf.register_new_resolver("eval", eval, replace=True) + @hydra.main(config_path=abs_config_path, version_base=None) def main(cfg): OmegaConf.resolve(cfg) @@ -145,6 +146,7 @@ def main(cfg): runner: BaseRunner = cls(cfg) runner.run() + if __name__ == "__main__": main() ``` diff --git a/docs/source/roboverse_learn/reinforcement_learning/fast_td3.md b/docs/source/roboverse_learn/reinforcement_learning/fast_td3.md index d258324ef..a60bc1e94 100644 --- a/docs/source/roboverse_learn/reinforcement_learning/fast_td3.md +++ b/docs/source/roboverse_learn/reinforcement_learning/fast_td3.md @@ -13,45 +13,39 @@ The script uses a `CONFIG` dictionary for all parameters. Key options include: ```python CONFIG = { # Environment - "sim": "mjx", # Simulator backend - "robots": ["h1"], # Robot type - "task": "humanoid.run", # Task name - "num_envs": 1024, # Number of parallel environments - "decimation": 10, # Control decimation - + "sim": "mjx", # Simulator backend + "robots": ["h1"], # Robot type + "task": "humanoid.run", # Task name + "num_envs": 1024, # Number of parallel environments + "decimation": 10, # Control decimation # Training - "total_timesteps": 1500, # Total training steps - "batch_size": 32768, # Batch size for updates - "buffer_size": 20480, # Replay buffer size - + "total_timesteps": 1500, # Total training steps + "batch_size": 32768, # Batch size for updates + "buffer_size": 20480, # Replay buffer size # Algorithm - "gamma": 0.99, # Discount factor - "tau": 0.1, # Target network update rate - "policy_frequency": 2, # Policy update frequency - "num_updates": 12, # Updates per step - + "gamma": 0.99, # Discount factor + "tau": 0.1, # Target network update rate + "policy_frequency": 2, # Policy update frequency + "num_updates": 12, # Updates per step # Networks "critic_learning_rate": 0.0003, "actor_learning_rate": 0.0003, "critic_hidden_dim": 1024, "actor_hidden_dim": 512, - # Distributional Q-learning "num_atoms": 101, "v_min": -250.0, "v_max": 250.0, - # Optimizations - "use_cdq": True, # Clipped Double Q-learning - "compile": True, # PyTorch compilation - "obs_normalization": True, # Observation normalization - "amp": True, # Automatic mixed precision - "amp_dtype": "fp16", # Precision type - + "use_cdq": True, # Clipped Double Q-learning + "compile": True, # PyTorch compilation + "obs_normalization": True, # Observation normalization + "amp": True, # Automatic mixed precision + "amp_dtype": "fp16", # Precision type # Logging - "use_wandb": False, # Weights & Biases integration - "eval_interval": 700, # Evaluation frequency - "save_interval": 700, # Model saving frequency + "use_wandb": False, # Weights & Biases integration + "eval_interval": 700, # Evaluation frequency + "save_interval": 700, # Model saving frequency } ``` diff --git a/docs/source/roboverse_learn/reinforcement_learning/humanoid.md b/docs/source/roboverse_learn/reinforcement_learning/humanoid.md index 171638cc3..dba785caa 100644 --- a/docs/source/roboverse_learn/reinforcement_learning/humanoid.md +++ b/docs/source/roboverse_learn/reinforcement_learning/humanoid.md @@ -264,8 +264,8 @@ ground_cfg = GroundCfg( dynamic_friction=1.0, elements={ "slope": [SlopeCfg(origin=[0, 0], size=[2.0, 2.0], slope=0.3)], - "stair": [StairCfg(origin=[5, 0], size=[2.0, 2.0], step_height=0.1)] - } + "stair": [StairCfg(origin=[5, 0], size=[2.0, 2.0], step_height=0.1)], + }, ) ``` diff --git a/examples/obj_layout/README.md b/examples/obj_layout/README.md index 982d75d53..f6266451e 100644 --- a/examples/obj_layout/README.md +++ b/examples/obj_layout/README.md @@ -119,14 +119,14 @@ class PutBananaTask(Task): name="banana", usd_path="...", pos=(0.5, 0.2, 0.15), # ← From saved poses - rot=(0, 0, 0, 1), # ← From saved poses + rot=(0, 0, 0, 1), # ← From saved poses ), ], robots=[ RobotCfg( name="franka", - pos=(0.0, 0.0, 0.0), # ← From saved poses - rot=(0, 0, 0, 1), # ← From saved poses + pos=(0.0, 0.0, 0.0), # ← From saved poses + rot=(0, 0, 0, 1), # ← From saved poses ), ], ) diff --git a/examples/rl/fast_td3/fttd3_module.py b/examples/rl/fast_td3/fttd3_module.py index fbb71789e..3b4b651b9 100644 --- a/examples/rl/fast_td3/fttd3_module.py +++ b/examples/rl/fast_td3/fttd3_module.py @@ -431,7 +431,8 @@ def projection( next_dist = F.softmax(self.forward(obs, actions), dim=1) proj_dist = torch.zeros_like(next_dist) offset = ( - torch.linspace(0, (batch_size - 1) * self.num_atoms, batch_size, device=device) + torch + .linspace(0, (batch_size - 1) * self.num_atoms, batch_size, device=device) .unsqueeze(1) .expand(batch_size, self.num_atoms) .long() diff --git a/examples/viser/VISER_USAGE_GUIDE.md b/examples/viser/VISER_USAGE_GUIDE.md index 728a5df2d..9052000b2 100644 --- a/examples/viser/VISER_USAGE_GUIDE.md +++ b/examples/viser/VISER_USAGE_GUIDE.md @@ -185,8 +185,8 @@ visualizer.enable_trajectory_playback() visualizer.setup_ik_solver(robot_name, robot_config, env_handler) # Load and set trajectory -visualizer.load_trajectory('trajectory.pkl.gz') -visualizer.set_current_trajectory('franka', 0) +visualizer.load_trajectory("trajectory.pkl.gz") +visualizer.set_current_trajectory("franka", 0) ``` ### Custom Scene Setup diff --git a/packages/metasim/docs/source/concept/config.md b/packages/metasim/docs/source/concept/config.md index 838e895f0..a7eaedccb 100644 --- a/packages/metasim/docs/source/concept/config.md +++ b/packages/metasim/docs/source/concept/config.md @@ -33,10 +33,8 @@ class ScenarioCfg: # Runtime render: RenderCfg = RenderCfg() sim_params: SimParamCfg = SimParamCfg() - simulator: Literal["isaaclab","isaacgym","sapien2","sapien3", - "genesis","pybullet","mujoco"] | None = None - renderer: Literal["isaaclab","isaacgym","sapien2","sapien3", - "genesis","pybullet","mujoco"] | None = None + simulator: Literal["isaaclab", "isaacgym", "sapien2", "sapien3", "genesis", "pybullet", "mujoco"] | None = None + renderer: Literal["isaaclab", "isaacgym", "sapien2", "sapien3", "genesis", "pybullet", "mujoco"] | None = None # Misc num_envs: int = 1 @@ -153,13 +151,12 @@ The class provides mechanisms for asset management and dynamic updates: ```python RobotCfg( - name="robot_template", - num_joints=2, - urdf_path="roboverse_data/robots/your_robot/urdf/your_robot.urdf", - fix_base_link=True, - enabled_gravity=True, - control_type={"joint1": "position", "joint2": "effort"}, - actuators={"joint1": BaseActuatorCfg(stiffness=500, damping=10), - "joint2": BaseActuatorCfg(effort_limit_sim=50)} + name="robot_template", + num_joints=2, + urdf_path="roboverse_data/robots/your_robot/urdf/your_robot.urdf", + fix_base_link=True, + enabled_gravity=True, + control_type={"joint1": "position", "joint2": "effort"}, + actuators={"joint1": BaseActuatorCfg(stiffness=500, damping=10), "joint2": BaseActuatorCfg(effort_limit_sim=50)}, ) ``` diff --git a/packages/metasim/docs/source/concept/get_extras.md b/packages/metasim/docs/source/concept/get_extras.md index 68ddc304e..78d7426b3 100644 --- a/packages/metasim/docs/source/concept/get_extras.md +++ b/packages/metasim/docs/source/concept/get_extras.md @@ -53,8 +53,8 @@ Implement the `_extra_spec` method inside your task definition: ```python def _extra_spec(self) -> dict: return { - "imu_pos": SitePos("imu"), # World position of IMU sensor site - "head_vel": BodyVel("head"), # Velocity of the robot's head + "imu_pos": SitePos("imu"), # World position of IMU sensor site + "head_vel": BodyVel("head"), # Velocity of the robot's head } ``` @@ -90,10 +90,7 @@ extras = handler.get_extra() This returns a dictionary containing your custom observations: ```python -{ - "imu_pos": tensor([[0.1, 0.2, 0.3], [...]]), - "head_vel": tensor([[0.01, 0.02, 0.03], [...]]) -} +{"imu_pos": tensor([[0.1, 0.2, 0.3], [...]]), "head_vel": tensor([[0.01, 0.02, 0.03], [...]])} ``` You can directly use this data for: @@ -134,6 +131,7 @@ Example: from metasim.queries.base import BaseQueryType import torch + class BodyMassQuery(BaseQueryType): def bind_handler(self, handler): # Store the handler for later use @@ -158,7 +156,7 @@ After defining your custom query class, simply instantiate and use it directly i ```python def _extra_spec(self) -> dict: return { - "body_mass": BodyMassQuery(), # Your custom query + "body_mass": BodyMassQuery(), # Your custom query "imu_pos": SitePos("imu"), } ``` diff --git a/packages/metasim/docs/source/concept/randomization.md b/packages/metasim/docs/source/concept/randomization.md index 0179babd6..51675a9cf 100644 --- a/packages/metasim/docs/source/concept/randomization.md +++ b/packages/metasim/docs/source/concept/randomization.md @@ -60,8 +60,8 @@ registry = ObjectRegistry.get_instance() # Query objects all_objects = registry.list_objects() -static_only = registry.list_objects(lifecycle='static') -dynamic_only = registry.list_objects(lifecycle='dynamic') +static_only = registry.list_objects(lifecycle="static") +dynamic_only = registry.list_objects(lifecycle="dynamic") # Get object metadata meta = registry.get_object_metadata("table") @@ -123,7 +123,7 @@ ManualGeometryCfg( geometry_type="cube", size=(10.0, 10.0, 0.1), position=(0.0, 0.0, 0.005), - default_material="roboverse_data/materials/arnold/Carpet/Carpet_Beige.mdl" + default_material="roboverse_data/materials/arnold/Carpet/Carpet_Beige.mdl", ) ``` @@ -135,7 +135,7 @@ USDAssetCfg( name="table", usd_path="EmbodiedGenData/dataset/basic_furniture/table/uuid.urdf", position=(0.0, 0.0, 0.0), - scale=(1.2, 1.5, 1.0) + scale=(1.2, 1.5, 1.0), ) ``` @@ -150,7 +150,7 @@ USDAssetPoolCfg( name="table_pool", usd_paths=[...], # Multiple USD files selection_strategy="random", - per_path_overrides={...} # Per-asset calibrations + per_path_overrides={...}, # Per-asset calibrations ) ``` @@ -169,8 +169,8 @@ MaterialRandomCfg( obj_name="box_base", mdl=MDLMaterialCfg( mdl_paths=[...], # Paths to MDL files - randomize_material_variant=True # Select from variants within files - ) + randomize_material_variant=True, # Select from variants within files + ), ) ``` @@ -183,11 +183,7 @@ Material application includes: **Physical Materials** (optional): ```python -PhysicalMaterialCfg( - friction_range=(0.3, 0.7), - restitution_range=(0.1, 0.3), - enabled=True -) +PhysicalMaterialCfg(friction_range=(0.3, 0.7), restitution_range=(0.1, 0.3), enabled=True) ``` Modifies friction and restitution on physics-enabled objects. Note that dynamic objects (created by SceneRandomizer) are visual-only and skip physical material randomization. @@ -209,16 +205,12 @@ Randomizes physics properties of static objects: ```python ObjectRandomCfg( obj_name="box_base", - physics=PhysicsRandomCfg( - mass_range=(10.0, 30.0), - friction_range=(0.3, 0.8), - enabled=True - ), + physics=PhysicsRandomCfg(mass_range=(10.0, 30.0), friction_range=(0.3, 0.8), enabled=True), pose=PoseRandomCfg( position_range=[(-0.1, 0.1), (-0.1, 0.1), (0, 0)], rotation_range=(0, 30), # Degrees - enabled=False # Disabled in trajectory replay - ) + enabled=False, # Disabled in trajectory replay + ), ) ``` @@ -231,25 +223,18 @@ Controls lighting parameters with support for multiple light types: ```python LightRandomCfg( light_name="ceiling_main", - intensity=LightIntensityRandomCfg( - intensity_range=(16000, 30000), - enabled=True - ), + intensity=LightIntensityRandomCfg(intensity_range=(16000, 30000), enabled=True), color=LightColorRandomCfg( temperature_range=(3000, 6000), # Kelvin use_temperature=True, - enabled=True + enabled=True, ), position=LightPositionRandomCfg( - position_range=((-1, 1), (-1, 1), (-0.2, 0.2)), - relative_to_origin=True, - enabled=True + position_range=((-1, 1), (-1, 1), (-0.2, 0.2)), relative_to_origin=True, enabled=True ), orientation=LightOrientationRandomCfg( - angle_range=((-20, 20), (-20, 20), (-180, 180)), - relative_to_origin=True, - enabled=True - ) + angle_range=((-20, 20), (-20, 20), (-180, 180)), relative_to_origin=True, enabled=True + ), ) ``` @@ -265,24 +250,13 @@ Perturbs camera parameters to simulate sensor variations: CameraRandomCfg( camera_name="main_camera", position=CameraPositionRandomCfg( - delta_range=((-0.1, 0.1), (-0.1, 0.1), (-0.05, 0.05)), - use_delta=True, - enabled=True - ), - orientation=CameraOrientationRandomCfg( - rotation_delta=((-5, 5), (-5, 5), (-5, 5)), - enabled=True + delta_range=((-0.1, 0.1), (-0.1, 0.1), (-0.05, 0.05)), use_delta=True, enabled=True ), + orientation=CameraOrientationRandomCfg(rotation_delta=((-5, 5), (-5, 5), (-5, 5)), enabled=True), look_at=CameraLookAtRandomCfg( - look_at_delta=((-0.05, 0.05), (-0.05, 0.05), (-0.05, 0.05)), - use_delta=True, - enabled=True + look_at_delta=((-0.05, 0.05), (-0.05, 0.05), (-0.05, 0.05)), use_delta=True, enabled=True ), - intrinsics=CameraIntrinsicsRandomCfg( - fov_range=(45, 60), - use_fov=True, - enabled=True - ) + intrinsics=CameraIntrinsicsRandomCfg(fov_range=(45, 60), use_fov=True, enabled=True), ) ``` @@ -349,10 +323,10 @@ This is controlled by `randomize_material_variant`: MDLMaterialCfg(mdl_paths=[...], randomize_material_variant=True) # Explicit variant specification -mdl_paths=["path/to/Wood.mdl::Oak"] +mdl_paths = ["path/to/Wood.mdl::Oak"] # Always use first variant in file -randomize_material_variant=False +randomize_material_variant = False ``` Variant selection is deterministic given a seed. The system parses MDL files to extract all available material names and uses the randomizer's RNG to select one. @@ -363,10 +337,7 @@ Variant selection is deterministic given a seed. The system parses MDL files to ```python # 1. Create randomizer with configuration -scene_cfg = ScenePresets.tabletop_workspace( - room_size=10.0, - wall_height=5.0 -) +scene_cfg = ScenePresets.tabletop_workspace(room_size=10.0, wall_height=5.0) scene_rand = SceneRandomizer(scene_cfg, seed=42) # 2. Bind to simulation handler @@ -412,8 +383,8 @@ randomizers = { "scene": SceneRandomizer(scene_cfg, seed=42), "material": MaterialRandomizer(mat_cfg, seed=43), "object": ObjectRandomizer(obj_cfg, seed=44), - "light": [LightRandomizer(light_cfg, seed=45+i) for i in range(5)], - "camera": CameraRandomizer(cam_cfg, seed=50) + "light": [LightRandomizer(light_cfg, seed=45 + i) for i in range(5)], + "camera": CameraRandomizer(cam_cfg, seed=50), } # Bind all @@ -424,6 +395,7 @@ for rand_list in randomizers.values(): else: rand.bind_handler(handler) + # Apply all def apply_all(level): randomizers["scene"]() @@ -513,7 +485,7 @@ ManualGeometryCfg( position=(0.0, 0.0, 0.65), rotation=(1.0, 0.0, 0.0, 0.0), # Quaternion (w, x, y, z) add_collision=True, # Add CollisionAPI for spatial queries - default_material="path/to/material.mdl" # Optional default appearance + default_material="path/to/material.mdl", # Optional default appearance ) ``` @@ -530,7 +502,7 @@ USDAssetCfg( position=(0.0, 0.0, 0.37), rotation=(1.0, 0.0, 0.0, 0.0), scale=(1.2, 1.5, 1.0), - auto_download=True + auto_download=True, ) ``` @@ -547,16 +519,12 @@ Asset pools enable geometric diversity: ```python USDAssetPoolCfg( name="table", - usd_paths=[ - "path/to/table1.urdf", - "path/to/table2.urdf", - "path/to/table3.urdf" - ], + usd_paths=["path/to/table1.urdf", "path/to/table2.urdf", "path/to/table3.urdf"], per_path_overrides={ "table1.urdf": {"position": (0, 0, 0.37), "scale": (1.2, 1.5, 1.0)}, - "table2.urdf": {"position": (0.3, 0, 0.37), "scale": (1.2, 1.4, 1.0)} + "table2.urdf": {"position": (0.3, 0, 0.37), "scale": (1.2, 1.4, 1.0)}, }, - selection_strategy="random" # or "sequential" + selection_strategy="random", # or "sequential" ) ``` @@ -582,10 +550,7 @@ To avoid consecutive repetition, the random strategy filters out the currently l from metasim.randomization import MaterialRandomizer, MaterialPresets # Using preset -mat_rand = MaterialRandomizer( - MaterialPresets.mdl_family_object("box_base", family=("wood", "metal")), - seed=123 -) +mat_rand = MaterialRandomizer(MaterialPresets.mdl_family_object("box_base", family=("wood", "metal")), seed=123) # Or manual configuration from metasim.randomization import MaterialRandomCfg, MDLMaterialCfg @@ -593,12 +558,9 @@ from metasim.randomization import MaterialRandomCfg, MDLMaterialCfg mat_rand = MaterialRandomizer( MaterialRandomCfg( obj_name="box_base", - mdl=MDLMaterialCfg( - mdl_paths=["path/to/Wood.mdl", "path/to/Metal.mdl"], - randomize_material_variant=True - ) + mdl=MDLMaterialCfg(mdl_paths=["path/to/Wood.mdl", "path/to/Metal.mdl"], randomize_material_variant=True), ), - seed=123 + seed=123, ) mat_rand.bind_handler(handler) @@ -621,10 +583,7 @@ A key feature of the refactored system is the ability to randomize materials on scene_rand() # Creates table with default Plywood material # Later: Randomize table material -table_mat = MaterialRandomizer( - MaterialPresets.mdl_family_object("table", family=("wood", "metal")), - seed=99 -) +table_mat = MaterialRandomizer(MaterialPresets.mdl_family_object("table", family=("wood", "metal")), seed=99) table_mat.bind_handler(handler) table_mat() # Changes table material to random wood or metal ``` @@ -644,6 +603,7 @@ MaterialPresets.mdl_family_object("floor", family=("carpet", "wood", "stone")) # Specific collection from metasim.randomization.presets import MDLCollections + paths = MDLCollections.family("metal") # All metal materials ``` @@ -679,7 +639,7 @@ wall_seed = base_seed + 100 for wall_name in ["wall_front", "wall_back", "wall_left", "wall_right"]: wall_mat = MaterialRandomizer( MaterialPresets.mdl_family_object(wall_name, family="masonry"), - seed=wall_seed # Same seed for all + seed=wall_seed, # Same seed for all ) wall_mat.bind_handler(handler) wall_mat() @@ -703,6 +663,7 @@ RoboVerse supports Hybrid mode, where IsaacLab manages physics and IsaacSim hand class SceneRandomizer(BaseRandomizerType): REQUIRES_HANDLER = "render" # Needs IsaacSim for USD operations + class ObjectRandomizer(BaseRandomizerType): REQUIRES_HANDLER = "physics" # Needs IsaacLab for physics APIs ``` @@ -718,7 +679,7 @@ For vectorized training with different randomization per environment: scene_cfg = SceneRandomCfg( workspace_layer=WorkspaceLayerCfg( shared=False, # Different table per env - elements=[USDAssetPoolCfg(...)] + elements=[USDAssetPoolCfg(...)], ) ) @@ -759,7 +720,7 @@ class MyPresets: ManualGeometryCfg( name="floor", size=(20.0, 20.0, 0.1), - default_material="roboverse_data/materials/arnold/Concrete/Concrete_Polished.mdl" + default_material="roboverse_data/materials/arnold/Concrete/Concrete_Polished.mdl", ), # ... walls, ceiling ] @@ -768,12 +729,13 @@ class MyPresets: elements=[ USDAssetPoolCfg( name="workbench", - usd_paths=[...] # Industrial tables + usd_paths=[...], # Industrial tables ) ] - ) + ), ) + # Use in task scene_rand = SceneRandomizer(MyPresets.industrial_workspace(), seed=42) ``` diff --git a/packages/metasim/docs/source/concept/state.md b/packages/metasim/docs/source/concept/state.md index 1c05884c0..dd90029d5 100644 --- a/packages/metasim/docs/source/concept/state.md +++ b/packages/metasim/docs/source/concept/state.md @@ -139,6 +139,7 @@ field the backend did not report is an empty list. ```python Dof = Dict[str, float] + class DictObjectState(TypedDict): """State of the object.""" @@ -160,13 +161,14 @@ class DictRobotState(DictObjectState): dof_vel_target: Dof | None dof_torque: Dof | None + class DictEnvState(TypedDict): """State of the environment.""" objects: dict[str, DictObjectState] robots: dict[str, DictRobotState] cameras: dict[str, dict[str, torch.Tensor]] - extras: dict[str, Any] # States of Extra information + extras: dict[str, Any] # States of Extra information ``` To obtain the `DictEnvState` from a handler, one can use `handler.get_state(mode="dict")` method. The return value will be a `DictEnvState` describing the current simulation status. @@ -213,21 +215,23 @@ Here's a concrete example of controlling a Franka Panda robot. The Franka has 7 ```python # Example 1: Set Franka to home position, num_envs=1 -actions = [{ - "franka": { - "dof_pos_target": { - "panda_joint1": 0.0, - "panda_joint2": -0.785398, - "panda_joint3": 0.0, - "panda_joint4": -2.356194, - "panda_joint5": 0.0, - "panda_joint6": 1.570796, - "panda_joint7": 0.785398, - "panda_finger_joint1": 0.04, - "panda_finger_joint2": 0.04, +actions = [ + { + "franka": { + "dof_pos_target": { + "panda_joint1": 0.0, + "panda_joint2": -0.785398, + "panda_joint3": 0.0, + "panda_joint4": -2.356194, + "panda_joint5": 0.0, + "panda_joint6": 1.570796, + "panda_joint7": 0.785398, + "panda_finger_joint1": 0.04, + "panda_finger_joint2": 0.04, + } } } -}] +] handler.set_dof_target(actions) ``` @@ -235,18 +239,19 @@ handler.set_dof_target(actions) # Example 2: Random joint positions within joint limits, num_envs=1 import torch -actions = [{ - "franka": { - "dof_pos_target": { - joint_name: ( - torch.rand(1).item() - * (robot.joint_limits[joint_name][1] - robot.joint_limits[joint_name][0]) - + robot.joint_limits[joint_name][0] - ) - for joint_name in robot.joint_limits.keys() +actions = [ + { + "franka": { + "dof_pos_target": { + joint_name: ( + torch.rand(1).item() * (robot.joint_limits[joint_name][1] - robot.joint_limits[joint_name][0]) + + robot.joint_limits[joint_name][0] + ) + for joint_name in robot.joint_limits.keys() + } } } -}] +] handler.set_dof_target(actions) ``` @@ -354,11 +359,7 @@ Where `DictTraj` is a dict with the following structure: ```python dict_traj = { "actions": List[Dict[str, np.ndarray]], # List of actions - "states": [ - DictTrajState1, - DictTrajState2, - ... - ] + "states": [DictTrajState1, DictTrajState2, ...], } ``` diff --git a/packages/metasim/docs/source/concept/task.md b/packages/metasim/docs/source/concept/task.md index 8c982da6d..41f967ba9 100644 --- a/packages/metasim/docs/source/concept/task.md +++ b/packages/metasim/docs/source/concept/task.md @@ -40,7 +40,6 @@ scenario = task_cls.scenario.update( # Create task env via registry device = torch.device("cuda" if torch.cuda.is_available() else "cpu") env = task_cls(scenario=scenario, device=device) - ``` **Key points:** @@ -65,8 +64,10 @@ This workflow ensures tasks are: ```python """Train PPO for a reaching task using RLTaskEnv.""" + from metasim.task.registry import get_task_class import torch + task_cls = get_task_class(args.task) # e.g., "example.my_task" # Start from the class-provided default scenario and override as needed @@ -93,7 +94,7 @@ from gymnasium import make_vec metasim.register_gym_envs() # register RoboVerse/* with Gymnasium before make_vec env = make_vec( - env_id, # e.g., "example.my_task" + env_id, # e.g., "example.my_task" num_envs=args.num_envs, robots=[args.robot], simulator=args.sim, @@ -157,9 +158,11 @@ from metasim.task.base import BaseTaskEnv from metasim.task.registry import register_task from metasim.scenario.scenario import ScenarioCfg + @register_task("example.my_task") class MyExampleTask(BaseTaskEnv): scenario = ScenarioCfg(robots=["franka"], simulator="mujoco", cameras=[]) + def _observation(self, state): ... def _privileged_observation(self, state): ... def _reward(self, state, action, next_state=None): ... @@ -168,8 +171,8 @@ class MyExampleTask(BaseTaskEnv): def _observation_space(self): ... def _action_space(self): ... def _extra_spec(self): ... - def step(self,actions): ... - def reset(self,states,env_ids): ... + def step(self, actions): ... + def reset(self, states, env_ids): ... ``` ### 4.3 Using the Task Template diff --git a/packages/metasim/docs/source/developer_guide/architecture_review.md b/packages/metasim/docs/source/developer_guide/architecture_review.md index aa853e268..52f92111b 100644 --- a/packages/metasim/docs/source/developer_guide/architecture_review.md +++ b/packages/metasim/docs/source/developer_guide/architecture_review.md @@ -211,10 +211,12 @@ clear failing test on day one if they forget a method. ```python # Method 1: Gymnasium style (used in clean_rl, vla) from gymnasium import make_vec + env = make_vec("RoboVerse/task", robots=[...], simulator=sim) # Method 2: Direct task class (used in fast_td3, rsl_rl, il) from metasim.task.registry import get_task_class + env = get_task_class(task)(scenario) ``` @@ -498,6 +500,7 @@ once. ```python # metasim/sim/base.py + class BaseSimHandler(ABC): def __init__(self, scenario, optional_queries=None): # ... existing code ... @@ -517,9 +520,7 @@ class BaseSimHandler(ABC): self._set_states(states, env_ids) def get_states( - self, - env_ids: list[int] | None = None, - mode: Literal["tensor", "dict"] = "tensor" + self, env_ids: list[int] | None = None, mode: Literal["tensor", "dict"] = "tensor" ) -> TensorState | list[DictEnvState]: """Get states with independent caching for each mode.""" if self._state_cache_expire: @@ -531,9 +532,7 @@ class BaseSimHandler(ABC): return self._tensor_state_cache else: if self._dict_state_cache is None: - self._dict_state_cache = state_tensor_to_nested( - self, self._tensor_state_cache - ) + self._dict_state_cache = state_tensor_to_nested(self, self._tensor_state_cache) return self._dict_state_cache ``` @@ -542,24 +541,25 @@ class BaseSimHandler(ABC): ```python # metasim/test/sim/test_state_cache.py + @pytest.mark.general def test_state_cache_mode_independence(): """Verify that switching modes doesn't corrupt cache.""" handler = create_test_handler() handler.launch() - + # Get tensor state states_t1 = handler.get_states(mode="tensor") assert isinstance(states_t1, TensorState) - + # Get dict state (should not affect tensor cache) states_d = handler.get_states(mode="dict") assert isinstance(states_d, list) - + # Get tensor state again (should return same type) states_t2 = handler.get_states(mode="tensor") assert isinstance(states_t2, TensorState) - + # Values should match assert torch.allclose(states_t1.pos, states_t2.pos) ``` @@ -588,6 +588,7 @@ def test_state_cache_mode_independence(): def _set_dof_targets(self, actions: list[Action]) -> None: raise NotImplementedError + # After @abstractmethod def _set_dof_targets(self, actions: list[Action]) -> None: @@ -619,27 +620,27 @@ from pathlib import Path import importlib import pkgutil + def get_all_robot_configs(): """Dynamically discover all robot configuration classes.""" import roboverse_pack.robots as robots_module - + configs = [] for importer, modname, ispkg in pkgutil.iter_modules(robots_module.__path__): - if modname.endswith('_cfg'): - module = importlib.import_module(f'roboverse_pack.robots.{modname}') + if modname.endswith("_cfg"): + module = importlib.import_module(f"roboverse_pack.robots.{modname}") for name in dir(module): obj = getattr(module, name) - if (isinstance(obj, type) and - hasattr(obj, 'name') and - name.endswith('Cfg')): + if isinstance(obj, type) and hasattr(obj, "name") and name.endswith("Cfg"): configs.append(obj) return configs + ALL_ROBOT_CONFIGS = get_all_robot_configs() + @pytest.mark.general -@pytest.mark.parametrize("robot_cfg_cls", ALL_ROBOT_CONFIGS, - ids=lambda x: x.__name__) +@pytest.mark.parametrize("robot_cfg_cls", ALL_ROBOT_CONFIGS, ids=lambda x: x.__name__) def test_robot_config_instantiation(robot_cfg_cls): """Verify robot config can be instantiated.""" cfg = robot_cfg_cls() @@ -647,43 +648,42 @@ def test_robot_config_instantiation(robot_cfg_cls): assert isinstance(cfg.name, str) assert len(cfg.name) > 0 + @pytest.mark.general -@pytest.mark.parametrize("robot_cfg_cls", ALL_ROBOT_CONFIGS, - ids=lambda x: x.__name__) +@pytest.mark.parametrize("robot_cfg_cls", ALL_ROBOT_CONFIGS, ids=lambda x: x.__name__) def test_robot_config_has_asset_path(robot_cfg_cls): """Verify robot config has at least one asset path.""" cfg = robot_cfg_cls() - + asset_paths = [ - getattr(cfg, 'usd_path', None), - getattr(cfg, 'urdf_path', None), - getattr(cfg, 'mjcf_path', None), + getattr(cfg, "usd_path", None), + getattr(cfg, "urdf_path", None), + getattr(cfg, "mjcf_path", None), ] - + valid_paths = [p for p in asset_paths if p is not None and len(p) > 0] assert len(valid_paths) > 0, f"{cfg.name} has no valid asset path" + @pytest.mark.general -@pytest.mark.parametrize("robot_cfg_cls", ALL_ROBOT_CONFIGS, - ids=lambda x: x.__name__) +@pytest.mark.parametrize("robot_cfg_cls", ALL_ROBOT_CONFIGS, ids=lambda x: x.__name__) def test_robot_config_has_actuators(robot_cfg_cls): """Verify robot config has actuator definitions.""" cfg = robot_cfg_cls() - - if hasattr(cfg, 'actuators'): + + if hasattr(cfg, "actuators"): assert len(cfg.actuators) > 0, f"{cfg.name} has no actuators defined" + @pytest.mark.general -@pytest.mark.parametrize("robot_cfg_cls", ALL_ROBOT_CONFIGS, - ids=lambda x: x.__name__) +@pytest.mark.parametrize("robot_cfg_cls", ALL_ROBOT_CONFIGS, ids=lambda x: x.__name__) def test_robot_config_joint_limits_valid(robot_cfg_cls): """Verify joint limits are valid (lower < upper).""" cfg = robot_cfg_cls() - - if hasattr(cfg, 'joint_limits'): + + if hasattr(cfg, "joint_limits"): for joint_name, (lower, upper) in cfg.joint_limits.items(): - assert lower < upper, \ - f"{cfg.name}.{joint_name}: lower ({lower}) >= upper ({upper})" + assert lower < upper, f"{cfg.name}.{joint_name}: lower ({lower}) >= upper ({upper})" ``` --- @@ -749,12 +749,13 @@ from metasim.task.registry import get_task_class, TASK_REGISTRY # Select representative tasks for testing CORE_TASKS = [ "pick_cube", - "place_cube", + "place_cube", "open_drawer", "close_drawer", "push_button", ] + @pytest.mark.mujoco @pytest.mark.parametrize("task_name", CORE_TASKS) def test_task_reset_step_mujoco(task_name): @@ -766,19 +767,19 @@ def test_task_reset_step_mujoco(task_name): num_envs=1, headless=True, ) - + env = task_cls(scenario, device="cpu") env.launch() - + try: # Test reset obs, info = env.reset() assert obs is not None - + # Test step action = env.action_space.sample() obs, reward, terminated, truncated, info = env.step(action) - + assert obs is not None assert isinstance(reward, (int, float, torch.Tensor)) assert isinstance(terminated, (bool, torch.Tensor)) @@ -804,6 +805,7 @@ import pytest import torch from roboverse_learn.il.policies.dp.ddpm_dit_image_policy import DDPMDiTImagePolicy + @pytest.fixture def sample_obs(): """Create sample observation for testing.""" @@ -812,6 +814,7 @@ def sample_obs(): "agent_pos": torch.randn(1, 7), } + def test_diffusion_policy_forward(): """Test forward pass of diffusion policy.""" policy = DDPMDiTImagePolicy( @@ -820,20 +823,21 @@ def test_diffusion_policy_forward(): horizon=16, # ... minimal config ) - + obs = sample_obs() action = policy.predict_action(obs) - + assert action.shape == (1, 16, 7) # (batch, horizon, action_dim) + def test_diffusion_policy_training_step(): """Test single training step.""" policy = DDPMDiTImagePolicy(...) optimizer = torch.optim.Adam(policy.parameters()) - + batch = create_training_batch() loss = policy.compute_loss(batch) - + assert loss.requires_grad loss.backward() optimizer.step() @@ -855,6 +859,7 @@ def test_diffusion_policy_training_step(): from typing import Union, List, Optional from metasim.scenario.robot import RobotCfg + def make_env( task: str, robots: Optional[List[Union[str, RobotCfg]]] = None, @@ -863,13 +868,13 @@ def make_env( headless: bool = True, device: str = "cuda", cameras: Optional[List] = None, - **kwargs + **kwargs, ): """ Unified environment factory for RoboVerse. - + This is the recommended way to create environments. - + Args: task: Task name (e.g., "pick_cube", "locomotion_walk") robots: List of robot names or RobotCfg instances @@ -879,10 +884,10 @@ def make_env( device: Device for tensor computations cameras: Camera configurations for observations **kwargs: Additional task-specific arguments - + Returns: BaseTaskEnv: Configured environment instance - + Example: >>> env = make_env( ... task="pick_cube", @@ -895,13 +900,13 @@ def make_env( """ from metasim.task.registry import get_task_class from metasim.utils.setup_util import get_robot - + # Resolve task class task_cls = get_task_class(task) - + # Build scenario from task default scenario = task_cls.scenario.copy() - + # Resolve robots if robots is not None: resolved_robots = [] @@ -911,7 +916,7 @@ def make_env( else: resolved_robots.append(robot) scenario.robots = resolved_robots - + # Apply overrides scenario.update( simulator=simulator, @@ -919,19 +924,20 @@ def make_env( headless=headless, cameras=cameras or [], ) - + # Create and return environment env = task_cls(scenario, device=device, **kwargs) env.launch() - + return env + # Also register with gymnasium for compatibility def register_gymnasium_envs(): """Register all RoboVerse tasks with Gymnasium.""" import gymnasium from metasim.task.registry import TASK_REGISTRY - + for task_name in TASK_REGISTRY: gymnasium.register( id=f"RoboVerse/{task_name}", @@ -1030,15 +1036,16 @@ For high-risk changes, use feature flags: ```python # metasim/config.py + class FeatureFlags: """Feature flags for gradual rollout of changes.""" - + # State cache v2 with independent tensor/dict caches USE_INDEPENDENT_STATE_CACHE = False - + # New unified environment factory USE_UNIFIED_ENV_FACTORY = False - + # Strict type checking in configs STRICT_CONFIG_VALIDATION = False ``` diff --git a/packages/metasim/docs/source/developer_guide/autotest.md b/packages/metasim/docs/source/developer_guide/autotest.md index 78b921756..8b437b743 100644 --- a/packages/metasim/docs/source/developer_guide/autotest.md +++ b/packages/metasim/docs/source/developer_guide/autotest.md @@ -121,6 +121,7 @@ from metasim.scenario.scenario import ScenarioCfg from metasim.test.conftest import register_shared_suite from roboverse_pack.robots.panda_cfg import PandaCfg + def get_manipulation_scenario(sim: str, num_envs: int) -> ScenarioCfg: """Build scenario for all manipulation tests. @@ -137,6 +138,7 @@ def get_manipulation_scenario(sim: str, num_envs: int) -> ScenarioCfg: # ... other config ) + # Register this scenario for ALL tests in the manipulation/ directory register_shared_suite("metasim.test.manipulation", get_manipulation_scenario) ``` @@ -145,12 +147,14 @@ register_shared_suite("metasim.test.manipulation", get_manipulation_scenario) ```python import pytest + @pytest.mark.mujoco def test_gripper_open_close(handler): """Test gripper opening and closing.""" # handler is the shared handler instance for this (sim, num_envs) combo assert handler.scenario.simulator == "mujoco" + @pytest.mark.sim("mujoco", "isaacsim") def test_gripper_force_limits(handler): """Test gripper force limits on multiple backends.""" @@ -161,6 +165,7 @@ def test_gripper_force_limits(handler): ```python import pytest + @pytest.mark.isaacsim def test_grasp_cube(handler): """Test grasping a cube object.""" @@ -188,6 +193,7 @@ from metasim.test.conftest import register_shared_suite from roboverse_pack.robots.bipedal_cfg import BipedalCfg from roboverse_pack.robots.quadruped_cfg import QuadrupedCfg + def get_walking_scenario(sim: str, num_envs: int) -> ScenarioCfg: """Build scenario for bipedal walking tests.""" return ScenarioCfg( @@ -197,6 +203,7 @@ def get_walking_scenario(sim: str, num_envs: int) -> ScenarioCfg: headless=True, ) + def get_running_scenario(sim: str, num_envs: int) -> ScenarioCfg: """Build scenario for quadruped running tests.""" return ScenarioCfg( @@ -206,6 +213,7 @@ def get_running_scenario(sim: str, num_envs: int) -> ScenarioCfg: headless=True, ) + # Register TWO scenarios with file-specific prefixes register_shared_suite("metasim.test.locomotion.test_walking", get_walking_scenario) register_shared_suite("metasim.test.locomotion.test_running", get_running_scenario) @@ -215,12 +223,14 @@ register_shared_suite("metasim.test.locomotion.test_running", get_running_scenar ```python import pytest + @pytest.mark.mujoco def test_bipedal_gait(handler): """Test bipedal walking gait.""" # Uses BipedalCfg scenario assert len(handler.scenario.robots) == 1 + @pytest.mark.isaacsim def test_balance_control(handler): """Test balance during walking.""" @@ -232,12 +242,14 @@ def test_balance_control(handler): ```python import pytest + @pytest.mark.mujoco def test_quadruped_trot(handler): """Test quadruped trotting gait.""" # Uses QuadrupedCfg scenario (different from walking!) assert len(handler.scenario.robots) == 1 + @pytest.mark.isaacsim def test_high_speed_stability(handler): """Test stability at high running speeds.""" @@ -256,6 +268,7 @@ Tests marked `@pytest.mark.general` are pure unit tests that don't require any s import pytest from metasim.utils import some_pure_function + @pytest.mark.general def test_pure_math(): """Test a pure function that doesn't need any simulator.""" @@ -263,6 +276,7 @@ def test_pure_math(): result = some_pure_function(2, 3) assert result == 5 + @pytest.mark.general def test_config_validation(): """Test configuration validation logic.""" @@ -274,6 +288,7 @@ def test_config_validation(): assert cfg.num_envs == 1 assert cfg.simulator == "mujoco" + @pytest.mark.general def test_string_parsing(): """Test string parsing utilities.""" diff --git a/packages/metasim/docs/source/developer_guide/docstring.md b/packages/metasim/docs/source/developer_guide/docstring.md index 5097dd290..071b89d94 100644 --- a/packages/metasim/docs/source/developer_guide/docstring.md +++ b/packages/metasim/docs/source/developer_guide/docstring.md @@ -61,6 +61,7 @@ class RigidObjMetaCfg(BaseObjMetaCfg): mesh_path: str | None = None physics: PhysicStateType = MISSING + @configclass class PrimitiveCubeMetaCfg(RigidObjMetaCfg): """Primitive cube object cfg. @@ -86,7 +87,6 @@ class PrimitiveCubeMetaCfg(RigidObjMetaCfg): def density(self) -> float: """Object density, for SAPIEN usage.""" return self.mass / (self.size[0] * self.size[1] * self.size[2]) - ``` diff --git a/packages/metasim/docs/source/developer_guide/tips/index.md b/packages/metasim/docs/source/developer_guide/tips/index.md index 9904a51f5..c8dc0580b 100644 --- a/packages/metasim/docs/source/developer_guide/tips/index.md +++ b/packages/metasim/docs/source/developer_guide/tips/index.md @@ -32,6 +32,7 @@ pip install -e ".[dev]" ```python # Enable verbose logging import logging + logging.basicConfig(level=logging.DEBUG) # Run in non-headless mode to visualize diff --git a/packages/metasim/docs/source/features/hybrid_rendering.md b/packages/metasim/docs/source/features/hybrid_rendering.md index 8902409ce..ddbc7d254 100644 --- a/packages/metasim/docs/source/features/hybrid_rendering.md +++ b/packages/metasim/docs/source/features/hybrid_rendering.md @@ -37,15 +37,15 @@ from metasim.sim import HybridSimHandler from metasim.utils.setup_util import get_sim_handler_class from metasim.constants import SimType -scenario = ... # cameras defined here +scenario = ... # cameras defined here physics = get_sim_handler_class(SimType.MUJOCO)(scenario.replace(simulator="mujoco", cameras=[], headless=True)) renderer = get_sim_handler_class(SimType.ISAACSIM)(scenario.replace(simulator="isaacsim")) handler = HybridSimHandler(scenario, physics, renderer) -handler.launch() # starts Kit for the renderer +handler.launch() # starts Kit for the renderer handler.set_dof_targets(actions) handler.simulate() -state = handler.get_states(mode="tensor") # state.cameras["cam"].rgb: (N, H, W, 3) uint8 +state = handler.get_states(mode="tensor") # state.cameras["cam"].rgb: (N, H, W, 3) uint8 ``` If the process already hosts an Isaac Sim application (tests, notebooks, an outer `AppLauncher`), diff --git a/packages/metasim/docs/source/features/replayable_episodes.md b/packages/metasim/docs/source/features/replayable_episodes.md index 7fb324741..e0e9c7a3d 100644 --- a/packages/metasim/docs/source/features/replayable_episodes.md +++ b/packages/metasim/docs/source/features/replayable_episodes.md @@ -29,9 +29,9 @@ from metasim.utils.replay import verify_episode_replay episode = record_episode(handler, handler.get_states(mode="tensor"), actions, seed=0) save_episode(episode, "episode.npz") -back = load_episode("episode.npz") # validates format, shapes, names, quaternions -check_assets(back) # {"franka.mjcf_path": "ok" | "changed" | "missing"} -print(verify_episode_replay(handler, back)) # L0 from disk: replay the actions, compare every state +back = load_episode("episode.npz") # validates format, shapes, names, quaternions +check_assets(back) # {"franka.mjcf_path": "ok" | "changed" | "missing"} +print(verify_episode_replay(handler, back)) # L0 from disk: replay the actions, compare every state ``` `verify_episode_replay` refuses to compare across a different simulator, a different `num_envs` (a diff --git a/packages/metasim/docs/source/get_started/advanced/obj_layout/index.md b/packages/metasim/docs/source/get_started/advanced/obj_layout/index.md index c6a3ba4b7..8c1e1941c 100644 --- a/packages/metasim/docs/source/get_started/advanced/obj_layout/index.md +++ b/packages/metasim/docs/source/get_started/advanced/obj_layout/index.md @@ -119,14 +119,14 @@ class PutBananaTask(Task): name="banana", usd_path="...", pos=(0.5, 0.2, 0.15), # ← From saved poses - rot=(0, 0, 0, 1), # ← From saved poses + rot=(0, 0, 0, 1), # ← From saved poses ), ], robots=[ RobotCfg( name="franka", - pos=(0.0, 0.0, 0.0), # ← From saved poses - rot=(0, 0, 0, 1), # ← From saved poses + pos=(0.0, 0.0, 0.0), # ← From saved poses + rot=(0, 0, 0, 1), # ← From saved poses ), ], ) diff --git a/packages/metasim/docs/source/get_started/advanced/rl_example/infrastructure.md b/packages/metasim/docs/source/get_started/advanced/rl_example/infrastructure.md index fdad63da4..920fbb9b4 100644 --- a/packages/metasim/docs/source/get_started/advanced/rl_example/infrastructure.md +++ b/packages/metasim/docs/source/get_started/advanced/rl_example/infrastructure.md @@ -85,14 +85,14 @@ python examples/rl/1_fttd3.py ```python CONFIG = { - "sim": "mjx", # Simulator backend - "robots": ["h1"], # Robot type - "task": "humanoid.run", # Task name - "num_envs": 1024, # Number of environments - "total_timesteps": 1500, # Training timesteps - "batch_size": 32768, # Batch size - "learning_rate": 0.0003, # Learning rate - "use_wandb": False, # Weights & Biases logging + "sim": "mjx", # Simulator backend + "robots": ["h1"], # Robot type + "task": "humanoid.run", # Task name + "num_envs": 1024, # Number of environments + "total_timesteps": 1500, # Training timesteps + "batch_size": 32768, # Batch size + "learning_rate": 0.0003, # Learning rate + "use_wandb": False, # Weights & Biases logging } ``` diff --git a/packages/metasim/docs/source/get_started/quick_start/0_static_scene.md b/packages/metasim/docs/source/get_started/quick_start/0_static_scene.md index 7484c17bd..d5924c3c8 100644 --- a/packages/metasim/docs/source/get_started/quick_start/0_static_scene.md +++ b/packages/metasim/docs/source/get_started/quick_start/0_static_scene.md @@ -85,10 +85,12 @@ You will get the following image: **Initial State Setup**: Use `handler.set_states()` to position objects: ```python -init_states = [{ - "objects": { - "cube": {"pos": torch.tensor([0.3, -0.2, 0.05]), "rot": torch.tensor([1.0, 0.0, 0.0, 0.0])}, - "sphere": {"pos": torch.tensor([0.4, -0.6, 0.05]), "rot": torch.tensor([1.0, 0.0, 0.0, 0.0])}, +init_states = [ + { + "objects": { + "cube": {"pos": torch.tensor([0.3, -0.2, 0.05]), "rot": torch.tensor([1.0, 0.0, 0.0, 0.0])}, + "sphere": {"pos": torch.tensor([0.4, -0.6, 0.05]), "rot": torch.tensor([1.0, 0.0, 0.0, 0.0])}, + } } -}] +] ``` diff --git a/packages/metasim/docs/source/get_started/quick_start/12_domain_randomization.md b/packages/metasim/docs/source/get_started/quick_start/12_domain_randomization.md index 5db94a1c6..95ca10f6d 100644 --- a/packages/metasim/docs/source/get_started/quick_start/12_domain_randomization.md +++ b/packages/metasim/docs/source/get_started/quick_start/12_domain_randomization.md @@ -237,8 +237,8 @@ Good starting point for understanding how things work. Everything is just cube p **Environment**: ```python ScenePresets.empty_room( - room_size=10.0, # 10m x 10m room - wall_height=5.0, # 5m tall walls + room_size=10.0, # 10m x 10m room + wall_height=5.0, # 5m tall walls ) ``` @@ -254,7 +254,7 @@ ManualGeometryCfg( geometry_type="cube", size=(1.8, 1.8, 0.1), position=(0.0, 0.0, 0.65), - default_material="roboverse_data/materials/arnold/Wood/Plywood.mdl" + default_material="roboverse_data/materials/arnold/Wood/Plywood.mdl", ) ``` @@ -270,7 +270,7 @@ USDAssetPoolCfg( name="table", usd_paths=table_paths, # 5 tables from EmbodiedGen per_path_overrides=table_configs, # Per-table calibrations - selection_strategy="random" if level >= 1 else "sequential" + selection_strategy="random" if level >= 1 else "sequential", ) ``` @@ -286,7 +286,7 @@ USDAssetPoolCfg( name="kujiale_scene", usd_paths=scene_paths, # 12 interior scenes per_path_overrides=scene_configs, # Position/scale calibrations - selection_strategy="random" if level >= 1 else "sequential" + selection_strategy="random" if level >= 1 else "sequential", ) ``` @@ -307,9 +307,9 @@ Adds desktop objects to Mode 2 for maximum visual diversity. ObjectsLayerCfg( elements=[ USDAssetPoolCfg( - name=f"desktop_object_{i+1}", + name=f"desktop_object_{i + 1}", usd_paths=object_paths, # 10 fruit models - selection_strategy="random" + selection_strategy="random", ) for i in range(3) # Place 3 objects ] @@ -361,11 +361,7 @@ Selects from MDL families. Each MDL file may contain multiple variants. The syst **Physical Materials** (optional): ```python -PhysicalMaterialCfg( - friction_range=(0.3, 0.7), - restitution_range=(0.1, 0.3), - enabled=True -) +PhysicalMaterialCfg(friction_range=(0.3, 0.7), restitution_range=(0.1, 0.3), enabled=True) ``` Randomizes physics properties on objects with RigidBodyAPI. Note: dynamic objects from SceneRandomizer are visual-only, so they skip this. @@ -394,13 +390,13 @@ LightRandomCfg( light_name="ceiling_main", intensity=LightIntensityRandomCfg( intensity_range=(16000, 30000), # Raytracing - enabled=True + enabled=True, ), color=LightColorRandomCfg( temperature_range=(3000, 6000), # Warm to cool white use_temperature=True, - enabled=True - ) + enabled=True, + ), ) ``` @@ -495,9 +491,9 @@ After scene creation, the system adjusts robot and object positions to match the ```python table_bounds = scene_rand.get_table_bounds(env_id=0) if table_bounds: - table_height = table_bounds['height'] + table_height = table_bounds["height"] clearance = 0.05 - + for obj_state in init_state["objects"].values(): obj_state["pos"][2] = table_height + clearance for robot_state in init_state["robots"].values(): @@ -536,9 +532,9 @@ Each randomizer maintains an independent random number generator. Running the sc box_mat = MaterialRandomizer( MaterialPresets.mdl_family_object( "box_base", - family=("stone", "ceramic", "plastic") # Different families + family=("stone", "ceramic", "plastic"), # Different families ), - seed=args.seed + 1 + seed=args.seed + 1, ) ``` @@ -548,10 +544,9 @@ box_mat = MaterialRandomizer( # Add a fill light fill_light = LightRandomizer( LightRandomCfg( - light_name="fill_light", - intensity=LightIntensityRandomCfg(intensity_range=(5000, 10000), enabled=True) + light_name="fill_light", intensity=LightIntensityRandomCfg(intensity_range=(5000, 10000), enabled=True) ), - seed=args.seed + 20 + seed=args.seed + 20, ) randomizers["light"].append(fill_light) ``` @@ -567,7 +562,7 @@ environment_layer = EnvironmentLayerCfg( geometry_type="cube", size=(15.0, 15.0, 0.1), position=(0.0, 0.0, 0.005), - default_material="roboverse_data/materials/arnold/Stone/Marble.mdl" + default_material="roboverse_data/materials/arnold/Stone/Marble.mdl", ), # Add custom walls, ceiling, etc. ] @@ -629,8 +624,8 @@ ObjectRegistry gives unified access to all sim objects, whether they're static ( ```python registry = ObjectRegistry.get_instance() all_objects = registry.list_objects() -static_only = registry.list_objects(lifecycle='static') -dynamic_only = registry.list_objects(lifecycle='dynamic') +static_only = registry.list_objects(lifecycle="static") +dynamic_only = registry.list_objects(lifecycle="dynamic") ``` Material, Object, Light, and Camera randomizers query the registry automatically to find their targets. No need for manual prim path specification in most cases. @@ -713,17 +708,19 @@ class MyRandomCfg: obj_name: str my_param_range: tuple[float, float] + # Implement randomizer class MyRandomizer(BaseRandomizerType): def __init__(self, cfg: MyRandomCfg, seed: int | None = None): super().__init__(seed=seed) self.cfg = cfg - + def __call__(self): obj = self.registry.get_object(self.cfg.obj_name) value = self.rng.uniform(*self.cfg.my_param_range) # Apply randomization... + # Use in demo my_rand = MyRandomizer(cfg, seed=args.seed + 100) my_rand.bind_handler(handler) @@ -737,13 +734,11 @@ randomizers["custom"].append(my_rand) class SceneUSDCollections: @staticmethod def my_custom_assets( - *, - indices: list[int] | None = None, - return_configs: bool = False + *, indices: list[int] | None = None, return_configs: bool = False ) -> list[str] | tuple[list[str], dict]: paths = [...] # Your asset paths configs = {...} # Optional per-asset configs - + if return_configs: return (paths, configs) return paths diff --git a/packages/metasim/docs/source/get_started/quick_start/14_real_asset.md b/packages/metasim/docs/source/get_started/quick_start/14_real_asset.md index 1abbd873a..90af886cc 100644 --- a/packages/metasim/docs/source/get_started/quick_start/14_real_asset.md +++ b/packages/metasim/docs/source/get_started/quick_start/14_real_asset.md @@ -75,7 +75,7 @@ snapshot_download( allow_patterns="demo_assets/*", ) -target_asset_type = AssetType.MJCF # or AssetType.USD +target_asset_type = AssetType.MJCF # or AssetType.USD urdf_paths = [ f"{data_dir}/demo_assets/remote_control/result/remote_control.urdf", diff --git a/packages/metasim/docs/source/get_started/quick_start/17_rerun_visualization.md b/packages/metasim/docs/source/get_started/quick_start/17_rerun_visualization.md index 12eef349e..faf9c5bfa 100644 --- a/packages/metasim/docs/source/get_started/quick_start/17_rerun_visualization.md +++ b/packages/metasim/docs/source/get_started/quick_start/17_rerun_visualization.md @@ -183,8 +183,8 @@ from metasim.utils.rerun.rerun_util import RerunVisualizer # 1. Initialize visualizer visualizer = RerunVisualizer( app_name="My Simulation", - spawn=True, # Auto-open viewer - save_path="my_sim.rrd" # Optional: save recording + spawn=True, # Auto-open viewer + save_path="my_sim.rrd", # Optional: save recording ) # 2. Add coordinate frame (optional) @@ -198,11 +198,11 @@ visualizer.visualize_scenario_items(scenario.robots, robot_states) for step in range(num_steps): # Set timeline position visualizer.set_time(step) - + # Run simulation handler.simulate() obs = handler.get_states(mode="tensor") - + # Extract and update states for name, state in robot_states.items(): visualizer.update_item_pose(name, state) diff --git a/packages/metasim/docs/source/get_started/quick_start/1_control_robot.md b/packages/metasim/docs/source/get_started/quick_start/1_control_robot.md index a4192e9f6..28a1cd339 100644 --- a/packages/metasim/docs/source/get_started/quick_start/1_control_robot.md +++ b/packages/metasim/docs/source/get_started/quick_start/1_control_robot.md @@ -127,8 +127,7 @@ actions = [ robot.name: { "dof_pos_target": { joint_name: ( - torch.rand(1).item() - * (robot.joint_limits[joint_name][1] - robot.joint_limits[joint_name][0]) + torch.rand(1).item() * (robot.joint_limits[joint_name][1] - robot.joint_limits[joint_name][0]) + robot.joint_limits[joint_name][0] ) for joint_name in robot.joint_limits.keys() diff --git a/packages/metasim/docs/source/get_started/quick_start/2_add_new_robot.md b/packages/metasim/docs/source/get_started/quick_start/2_add_new_robot.md index b22e8afc0..a004fd02b 100644 --- a/packages/metasim/docs/source/get_started/quick_start/2_add_new_robot.md +++ b/packages/metasim/docs/source/get_started/quick_start/2_add_new_robot.md @@ -128,7 +128,7 @@ robot = RobotCfg( name="new_robot_h1", num_joints=26, usd_path="roboverse_data/robots/h1/usd/h1.usd", - mjcf_path="roboverse_data/robots/h1/mjcf/h1.xml", + mjcf_path="roboverse_data/robots/h1/mjcf/h1.xml", urdf_path="roboverse_data/robots/h1/urdf/h1.urdf", actuators={ "left_hip_yaw": BaseActuatorCfg(stiffness=200, damping=5), @@ -144,7 +144,7 @@ robot = RobotCfg( "left_hip_yaw": "position", "left_knee": "position", # ... more control modes - } + }, ) ``` diff --git a/packages/metasim/docs/source/get_started/quick_start/4_motion_planning.md b/packages/metasim/docs/source/get_started/quick_start/4_motion_planning.md index e5b58bc9b..fb8aed321 100644 --- a/packages/metasim/docs/source/get_started/quick_start/4_motion_planning.md +++ b/packages/metasim/docs/source/get_started/quick_start/4_motion_planning.md @@ -135,9 +135,9 @@ inverse_reorder_idx = [reorder_idx.index(i) for i in range(len(reorder_idx))] curr_robot_q = obs.robots[scenario.robots[0].name].joint_pos[:, inverse_reorder_idx] # Solve IK for multiple environments q_solution, ik_success = ik_solver.solve_ik_batch( - ee_pos_target=target_positions, # (B, 3) - target EE positions - ee_quat_target=target_quaternions, # (B, 4) - target EE quaternions (wxyz) - seed_q=curr_robot_q # (B, n_dof) - seed configs (required for curobo) + ee_pos_target=target_positions, # (B, 3) - target EE positions + ee_quat_target=target_quaternions, # (B, 4) - target EE quaternions (wxyz) + seed_q=curr_robot_q, # (B, n_dof) - seed configs (required for curobo) ) # q_solution: (B, n_dof_ik) - arm joint positions only @@ -154,15 +154,15 @@ This function combines the arm joint positions from IK with gripper positions to gripper_widths = process_gripper_command( gripper_binary=gripper_open_close, # (B,) or (B, 1) - binary gripper state robot_cfg=robot_cfg, - device=device + device=device, ) # Option 1: Return tensor in alphabetical order (default) actions_tensor = ik_solver.compose_joint_action( - q_solution=q_solution, # (B, n_dof_ik) - arm joint positions from IK - gripper_widths=gripper_widths, # (B, ee_n_dof) - gripper joint positions - current_q=current_joint_state, # (B, n_robot_dof) - optional current state - return_dict=False # Default: return tensor + q_solution=q_solution, # (B, n_dof_ik) - arm joint positions from IK + gripper_widths=gripper_widths, # (B, ee_n_dof) - gripper joint positions + current_q=current_joint_state, # (B, n_robot_dof) - optional current state + return_dict=False, # Default: return tensor ) # q_full: (B, n_robot_dof) - complete joint command in alphabetical order @@ -171,7 +171,7 @@ actions_dict = ik_solver.compose_joint_action( q_solution=q_solution, gripper_widths=gripper_widths, current_q=current_joint_state, - return_dict=True # Return action dictionaries + return_dict=True, # Return action dictionaries ) # actions: list of action dictionaries for env execution ``` @@ -213,25 +213,18 @@ inverse_reorder_idx = [reorder_idx.index(i) for i in range(len(reorder_idx))] curr_robot_q = obs.robots[scenario.robots[0].name].joint_pos[:, inverse_reorder_idx] q_arm, ik_success = ik_solver.solve_ik_batch( - ee_pos_target=target_positions, - ee_quat_target=target_quaternions, - seed_q=curr_robot_q + ee_pos_target=target_positions, ee_quat_target=target_quaternions, seed_q=curr_robot_q ) # Step 2: Process gripper commands -gripper_widths = process_gripper_command( - gripper_binary=gripper_commands, - robot_cfg=robot_cfg, - device=device -) +gripper_widths = process_gripper_command(gripper_binary=gripper_commands, robot_cfg=robot_cfg, device=device) # Step 3: Compose full joint command and create actions directly actions = ik_solver.compose_joint_action( q_solution=q_arm, gripper_widths=gripper_widths, - return_dict=True # Return action dictionaries directly + return_dict=True, # Return action dictionaries directly ) - ``` ### Backend-Specific Notes diff --git a/packages/metasim/docs/source/get_started/quick_start/guide.md b/packages/metasim/docs/source/get_started/quick_start/guide.md index 1a732050c..f87daa4d5 100644 --- a/packages/metasim/docs/source/get_started/quick_start/guide.md +++ b/packages/metasim/docs/source/get_started/quick_start/guide.md @@ -49,12 +49,7 @@ Here's how to create a basic scenario: from metasim.scenario.scenario import ScenarioCfg # Simple scenario with just a robot -scenario = ScenarioCfg( - robots=["franka"], - simulator="mujoco", - num_envs=1 -) - +scenario = ScenarioCfg(robots=["franka"], simulator="mujoco", num_envs=1) ``` Fields include: `robots` (robot configurations), `objects` (scene objects), `simulator` (physics engine), `num_envs` (number of parallel environments), `headless` (run without GUI), `cameras` (camera sensors), `lights` (lighting setup), and `sim_params` (physics parameters like timestep). @@ -89,23 +84,19 @@ from metasim.scenario.scenario import ScenarioCfg from metasim.utils.setup_util import get_handler # Define your scenario -scenario = ScenarioCfg( - robots=["franka"], - simulator="mujoco", - num_envs=1 -) +scenario = ScenarioCfg(robots=["franka"], simulator="mujoco", num_envs=1) # Create handler directly - one step handler = get_handler(scenario) # Control loop states = ... -handler.set_state(states) # Set state (for resets) +handler.set_state(states) # Set state (for resets) actions = ... # Your control actions -handler.set_dof_targets(actions) # Apply actions to robot +handler.set_dof_targets(actions) # Apply actions to robot -handler.simulate() # Step physics +handler.simulate() # Step physics obs = handler.get_state(mode="tensor") # Get updated state ``` @@ -127,18 +118,12 @@ task_cls = get_task_class(args.task) # e.g., "example.reaching" # Get default scenario and update with specific parameters scenario = task_cls.scenario.update( - robots=[args.robot], - simulator=args.sim, - num_envs=args.num_envs, - headless=args.headless, - cameras=[] + robots=[args.robot], simulator=args.sim, num_envs=args.num_envs, headless=args.headless, cameras=[] ) # Create task environment device = torch.device("cuda" if torch.cuda.is_available() else "cpu") env = task_cls(scenario=scenario, device=device) - - ``` ```python diff --git a/packages/metasim/metasim/randomization/core/isaacsim_adapter.py b/packages/metasim/metasim/randomization/core/isaacsim_adapter.py index bee53c720..ae3c3fcdf 100644 --- a/packages/metasim/metasim/randomization/core/isaacsim_adapter.py +++ b/packages/metasim/metasim/randomization/core/isaacsim_adapter.py @@ -263,7 +263,7 @@ def apply_mdl_material( # - material_name: Variant within MDL (optional, e.g., "Plywood::Oak") # - mtl_prim_name: Unique prim name to avoid conflicts mdl_basename = os.path.basename(mdl_path).removesuffix(".mdl") - prim_basename = prim_path.split("/")[-1] + prim_basename = prim_path.rsplit("/", maxsplit=1)[-1] mtl_prim_unique = f"{mdl_basename}_{prim_basename}_{id(prim_path)}" # Fully unique _, mtl_prim_path = get_material_prim_path(mtl_prim_unique) diff --git a/packages/metasim/metasim/sim/newton/newton.py b/packages/metasim/metasim/sim/newton/newton.py index 3dbfbd062..ca63d7e03 100644 --- a/packages/metasim/metasim/sim/newton/newton.py +++ b/packages/metasim/metasim/sim/newton/newton.py @@ -1574,7 +1574,8 @@ def _render_tiled_cameras(self, env_ids: list[int]) -> dict: depth_tensor = depth_tensor[use_env_ids] intrinsics = ( - torch.tensor(cam_cfg.intrinsics, dtype=torch.float32, device=self.device) + torch + .tensor(cam_cfg.intrinsics, dtype=torch.float32, device=self.device) .unsqueeze(0) .expand(len(use_env_ids), -1, -1) ) diff --git a/packages/metasim/metasim/utils/kinematics.py b/packages/metasim/metasim/utils/kinematics.py index 66584c191..f4e2eabbb 100644 --- a/packages/metasim/metasim/utils/kinematics.py +++ b/packages/metasim/metasim/utils/kinematics.py @@ -103,7 +103,8 @@ def _solve_ik_jax_with_seed(robot, target_link_index, target_wxyz, target_positi ] sol = ( - jaxls.LeastSquaresProblem(factors, [joint_var]) + jaxls + .LeastSquaresProblem(factors, [joint_var]) .analyze() .solve( initial_vals=jaxls.VarValues.make([joint_var.with_value(prev_cfg)]), @@ -130,7 +131,8 @@ def _solve_ik_jax_without_seed(robot, target_link_index, target_wxyz, target_pos ] sol = ( - jaxls.LeastSquaresProblem(factors, [joint_var]) + jaxls + .LeastSquaresProblem(factors, [joint_var]) .analyze() .solve( verbose=False, diff --git a/packages/metasim/metasim/utils/rerun/README.md b/packages/metasim/metasim/utils/rerun/README.md index 5809c217f..df8294075 100644 --- a/packages/metasim/metasim/utils/rerun/README.md +++ b/packages/metasim/metasim/utils/rerun/README.md @@ -19,16 +19,10 @@ from metasim.task.registry import get_task_class from metasim.utils.rerun.rerun_env_wrapper import TaskRerunWrapper # Create environment -task_cls = get_task_class('reach_origin') -scenario = task_cls.scenario.update( - robots=['franka'], - simulator='mujoco', - num_envs=1024, - headless=True, - cameras=[] -) +task_cls = get_task_class("reach_origin") +scenario = task_cls.scenario.update(robots=["franka"], simulator="mujoco", num_envs=1024, headless=True, cameras=[]) -env = task_cls(scenario, device='cuda') +env = task_cls(scenario, device="cuda") # Wrap environment to enable visualization rerun_env = TaskRerunWrapper(env, app_name="RL Training", update_freq=10) @@ -91,10 +85,10 @@ CONFIG = { ```python TaskRerunWrapper( - task_env, # RLTaskEnv or similar + task_env, # RLTaskEnv or similar app_name="RoboVerse Training", # Rerun application name - update_freq=10, # Update visualization every N steps - spawn=True, # Auto-open Rerun viewer + update_freq=10, # Update visualization every N steps + spawn=True, # Auto-open Rerun viewer ) ``` diff --git a/packages/metasim/metasim/utils/viser/README.md b/packages/metasim/metasim/utils/viser/README.md index 7e96aac6f..b2c788efe 100644 --- a/packages/metasim/metasim/utils/viser/README.md +++ b/packages/metasim/metasim/utils/viser/README.md @@ -19,16 +19,16 @@ from metasim.task.registry import get_task_class from metasim.utils.viser.viser_env_wrapper import TaskViserWrapper # 创建环境 -task_cls = get_task_class('reach_origin') +task_cls = get_task_class("reach_origin") scenario = task_cls.scenario.update( - robots=['franka'], - simulator='mujoco', + robots=["franka"], + simulator="mujoco", num_envs=1024, headless=False, # 启用渲染以支持 Viser - cameras=[] + cameras=[], ) -env = task_cls(scenario, device='cuda') +env = task_cls(scenario, device="cuda") # 包装环境以启用可视化 viser_env = TaskViserWrapper(env, port=8080) @@ -91,9 +91,9 @@ TaskViserWrapper 透明代理所有环境属性： ```python wrapper = TaskViserWrapper(env) -print(wrapper.num_envs) # 代理到 env.num_envs -print(wrapper.num_actions) # 代理到 env.num_actions -print(wrapper.num_obs) # 代理到 env.num_obs +print(wrapper.num_envs) # 代理到 env.num_envs +print(wrapper.num_actions) # 代理到 env.num_actions +print(wrapper.num_obs) # 代理到 env.num_obs print(wrapper.action_space) # 代理到 env.action_space ``` diff --git a/roboverse_learn/fusion/README.md b/roboverse_learn/fusion/README.md index 565e99a6a..cf35acb56 100644 --- a/roboverse_learn/fusion/README.md +++ b/roboverse_learn/fusion/README.md @@ -86,9 +86,9 @@ import torch from rsl_rl.modules import ActorCritic from roboverse_learn.fusion import load_bc_into_actor_critic -ac = ActorCritic(...) # fresh RL actor-critic -bc_sd = torch.load("bc_policy.pt") # a state-MLP BC policy -n = load_bc_into_actor_critic(ac, bc_sd) # actor now == BC; fine-tune with PPO +ac = ActorCritic(...) # fresh RL actor-critic +bc_sd = torch.load("bc_policy.pt") # a state-MLP BC policy +n = load_bc_into_actor_critic(ac, bc_sd) # actor now == BC; fine-tune with PPO ``` ## Tests diff --git a/roboverse_learn/il/policies/act/constants.py b/roboverse_learn/il/policies/act/constants.py index 012addaff..1fbe2756e 100644 --- a/roboverse_learn/il/policies/act/constants.py +++ b/roboverse_learn/il/policies/act/constants.py @@ -58,46 +58,46 @@ ############################ Helper functions ############################ -MASTER_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_POSITION_CLOSE) / ( - MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE +MASTER_GRIPPER_POSITION_NORMALIZE_FN = lambda x: ( + (x - MASTER_GRIPPER_POSITION_CLOSE) / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) ) -PUPPET_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_POSITION_CLOSE) / ( - PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE +PUPPET_GRIPPER_POSITION_NORMALIZE_FN = lambda x: ( + (x - PUPPET_GRIPPER_POSITION_CLOSE) / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) ) -MASTER_GRIPPER_POSITION_UNNORMALIZE_FN = ( - lambda x: x * (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) + MASTER_GRIPPER_POSITION_CLOSE +MASTER_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: ( + x * (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) + MASTER_GRIPPER_POSITION_CLOSE ) -PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN = ( - lambda x: x * (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + PUPPET_GRIPPER_POSITION_CLOSE +PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: ( + x * (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + PUPPET_GRIPPER_POSITION_CLOSE ) MASTER2PUPPET_POSITION_FN = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(MASTER_GRIPPER_POSITION_NORMALIZE_FN(x)) -MASTER_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_JOINT_CLOSE) / ( - MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE +MASTER_GRIPPER_JOINT_NORMALIZE_FN = lambda x: ( + (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) ) -PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_JOINT_CLOSE) / ( - PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE +PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: ( + (x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) ) -MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = ( - lambda x: x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE +MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: ( + x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE ) -PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = ( - lambda x: x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE +PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: ( + x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE ) MASTER2PUPPET_JOINT_FN = lambda x: PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(MASTER_GRIPPER_JOINT_NORMALIZE_FN(x)) MASTER_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) -MASTER_POS2JOINT = ( - lambda x: MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) +MASTER_POS2JOINT = lambda x: ( + MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE ) MASTER_JOINT2POS = lambda x: MASTER_GRIPPER_POSITION_UNNORMALIZE_FN( (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) ) -PUPPET_POS2JOINT = ( - lambda x: PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) +PUPPET_POS2JOINT = lambda x: ( + PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE ) PUPPET_JOINT2POS = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN( diff --git a/roboverse_learn/il/policies/act/detr/models/position_encoding.py b/roboverse_learn/il/policies/act/detr/models/position_encoding.py index 4292b8d79..17a123c49 100644 --- a/roboverse_learn/il/policies/act/detr/models/position_encoding.py +++ b/roboverse_learn/il/policies/act/detr/models/position_encoding.py @@ -87,7 +87,8 @@ def forward(self, tensor_list: NestedTensor): x_emb = self.col_embed(i) y_emb = self.row_embed(j) pos = ( - torch.cat( + torch + .cat( [ x_emb.unsqueeze(0).repeat(h, 1, 1), y_emb.unsqueeze(1).repeat(1, w, 1), diff --git a/roboverse_learn/rl/fast_td3/README.md b/roboverse_learn/rl/fast_td3/README.md index f3f429876..17aff639f 100644 --- a/roboverse_learn/rl/fast_td3/README.md +++ b/roboverse_learn/rl/fast_td3/README.md @@ -198,12 +198,11 @@ Add the following code in the standalone MetaSim checkout at `metasim/sim/isaacg ```python # eg RigidObjCfg after line 395 -asset_options.vhacd_enabled = True -asset_options.vhacd_params.resolution = 500000 # -asset_options.vhacd_params.max_convex_hulls = 64 # +asset_options.vhacd_enabled = True +asset_options.vhacd_params.resolution = 500000 # +asset_options.vhacd_params.max_convex_hulls = 64 # asset_options.vhacd_params.max_num_vertices_per_ch = 64 asset_options.thickness = 0.001 - ``` diff --git a/roboverse_learn/rl/fast_td3/fttd3_module.py b/roboverse_learn/rl/fast_td3/fttd3_module.py index 16ba660ec..0f8c54d9c 100644 --- a/roboverse_learn/rl/fast_td3/fttd3_module.py +++ b/roboverse_learn/rl/fast_td3/fttd3_module.py @@ -431,7 +431,8 @@ def projection( next_dist = F.softmax(self.forward(obs, actions), dim=1) proj_dist = torch.zeros_like(next_dist) offset = ( - torch.linspace(0, (batch_size - 1) * self.num_atoms, batch_size, device=device) + torch + .linspace(0, (batch_size - 1) * self.num_atoms, batch_size, device=device) .unsqueeze(1) .expand(batch_size, self.num_atoms) .long() diff --git a/roboverse_learn/vla/OpenVLA/vla_eval.py b/roboverse_learn/vla/OpenVLA/vla_eval.py index a50d4a497..742c4f53a 100644 --- a/roboverse_learn/vla/OpenVLA/vla_eval.py +++ b/roboverse_learn/vla/OpenVLA/vla_eval.py @@ -91,7 +91,8 @@ def _init_policy(self, **kwargs): self.processor = AutoProcessor.from_pretrained(self.model_path, trust_remote_code=True) self.model = ( - AutoModelForVision2Seq.from_pretrained( + AutoModelForVision2Seq + .from_pretrained( self.model_path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, diff --git a/roboverse_learn/vla/SmolVLA/smolvla_eval.py b/roboverse_learn/vla/SmolVLA/smolvla_eval.py index d160b4887..c8b52f7b0 100755 --- a/roboverse_learn/vla/SmolVLA/smolvla_eval.py +++ b/roboverse_learn/vla/SmolVLA/smolvla_eval.py @@ -113,7 +113,8 @@ def _init_policy(self, checkpoint_path: str): self.processor = AutoProcessor.from_pretrained(checkpoint_path, trust_remote_code=True) self.model = ( - AutoModel.from_pretrained( + AutoModel + .from_pretrained( checkpoint_path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, diff --git a/roboverse_learn/vla/pi0/README.md b/roboverse_learn/vla/pi0/README.md index 3cc5bd549..f2fe38933 100644 --- a/roboverse_learn/vla/pi0/README.md +++ b/roboverse_learn/vla/pi0/README.md @@ -150,15 +150,13 @@ Inside the openpi repo: def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig: repack_transform = _transforms.Group( inputs=[ - _transforms.RepackTransform( - { - "observation/image": "image", - # "observation/wrist_image": "wrist_image", # RoboVerse has a single view now - "observation/state": "state", - "actions": "actions", - "prompt": "prompt", - } - ) + _transforms.RepackTransform({ + "observation/image": "image", + # "observation/wrist_image": "wrist_image", # RoboVerse has a single view now + "observation/state": "state", + "actions": "actions", + "prompt": "prompt", + }) ] ) @@ -212,9 +210,7 @@ TrainConfig( decay_lr=5e-5, ), optimizer=_optimizer.AdamW(clip_gradient_norm=1.0), - weight_loader=weight_loaders.CheckpointWeightLoader( - "gs://openpi-assets/checkpoints/pi05_base/params" - ), + weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"), num_train_steps=30_000, freeze_filter=pi0_config.Pi0Config( paligemma_variant="gemma_2b_lora", diff --git a/roboverse_pack/callback_funcs/humanoid/step_funcs.py b/roboverse_pack/callback_funcs/humanoid/step_funcs.py index ee9c1ab95..d215a0060 100644 --- a/roboverse_pack/callback_funcs/humanoid/step_funcs.py +++ b/roboverse_pack/callback_funcs/humanoid/step_funcs.py @@ -80,7 +80,8 @@ def push_by_setting_velocity( / env.step_dt ).to(torch.int) push_env_ids = ( - torch.logical_and(env._episode_steps % env.push_interval == 0, env._episode_steps > 0) + torch + .logical_and(env._episode_steps % env.push_interval == 0, env._episode_steps > 0) .nonzero(as_tuple=False) .flatten() ) diff --git a/roboverse_pack/tasks/beyondmimic/metasim/envs/base_legged_robot.py b/roboverse_pack/tasks/beyondmimic/metasim/envs/base_legged_robot.py index 64e9f8db3..4cbdc91f6 100644 --- a/roboverse_pack/tasks/beyondmimic/metasim/envs/base_legged_robot.py +++ b/roboverse_pack/tasks/beyondmimic/metasim/envs/base_legged_robot.py @@ -138,7 +138,8 @@ def _instantiate_cfg(self): # per-actuator action scale self.action_scale = ( - torch.tensor(list(self.robot.action_scale.values()), device=self.device) + torch + .tensor(list(self.robot.action_scale.values()), device=self.device) .unsqueeze(0) .repeat(self.num_envs, 1) ) diff --git a/roboverse_pack/tasks/beyondmimic/metasim/mdp/events.py b/roboverse_pack/tasks/beyondmimic/metasim/mdp/events.py index 2bd9418e6..2b2bebf75 100644 --- a/roboverse_pack/tasks/beyondmimic/metasim/mdp/events.py +++ b/roboverse_pack/tasks/beyondmimic/metasim/mdp/events.py @@ -73,7 +73,8 @@ def push_by_setting_velocity( ).to(torch.int) # TODO different from how interval events are triggered in Isaac Lab, consider adapting this push_env_ids = ( - torch.logical_and(env._episode_steps % env.push_interval == 0, env._episode_steps > 0) + torch + .logical_and(env._episode_steps % env.push_interval == 0, env._episode_steps > 0) .nonzero(as_tuple=False) .flatten() ) diff --git a/roboverse_pack/tasks/beyondmimic/metasim/mdp/rewards.py b/roboverse_pack/tasks/beyondmimic/metasim/mdp/rewards.py index 5300b0bc4..2ff0c09b1 100644 --- a/roboverse_pack/tasks/beyondmimic/metasim/mdp/rewards.py +++ b/roboverse_pack/tasks/beyondmimic/metasim/mdp/rewards.py @@ -116,9 +116,8 @@ def undesired_contacts( contact_forces: ContactForces = env_states.extras["contact_forces"][env.name] is_contact = ( # TODO check correspondence with `contact_sensor.data.net_forces_w_history` - contact_forces.contact_forces_history[ - :, :, indexes, : - ] # [n_envs, history_length, n_indexes, 3] -> [n_envs, 3, 26, 3] + contact_forces + .contact_forces_history[:, :, indexes, :] # [n_envs, history_length, n_indexes, 3] -> [n_envs, 3, 26, 3] .norm(dim=-1) .max(dim=1)[0] > threshold diff --git a/roboverse_pack/tasks/box_task/README.md b/roboverse_pack/tasks/box_task/README.md index b33d71cd2..ce1c55d46 100644 --- a/roboverse_pack/tasks/box_task/README.md +++ b/roboverse_pack/tasks/box_task/README.md @@ -84,8 +84,11 @@ is missing locally, prepare it with the existing helper: ```python from roboverse_pack.asset.setup_interior_agent_assets import ( - download_roboverse_usd, copy_usd_assets, ensure_kujiale_scene, + download_roboverse_usd, + copy_usd_assets, + ensure_kujiale_scene, ) + download_roboverse_usd() copy_usd_assets(scene_ids=["0021", "0022", "0024", "0025", "0031"]) for sid in ["0021", "0022", "0024", "0025", "0031"]: diff --git a/roboverse_pack/tasks/mujoco_playground/NewTask.md b/roboverse_pack/tasks/mujoco_playground/NewTask.md index 111b31edf..ce4d05fc1 100644 --- a/roboverse_pack/tasks/mujoco_playground/NewTask.md +++ b/roboverse_pack/tasks/mujoco_playground/NewTask.md @@ -65,212 +65,194 @@ ```python def default_config() -> config_dict.ConfigDict: - """Returns the default config for bring_to_target tasks.""" - config = config_dict.create( - ctrl_dt=0.02, - sim_dt=0.005, - episode_length=150, - action_repeat=1, - action_scale=0.04, - reward_config=config_dict.create( - scales=config_dict.create( - # Gripper goes to the box. - gripper_box=4.0, - # Box goes to the target mocap. - box_target=8.0, - # Do not collide the gripper with the floor. - no_floor_collision=0.25, - # Arm stays close to target pose. - robot_target_qpos=0.3, - ) - ), - impl='jax', - nconmax=24 * 8192, - njmax=128, - ) - return config + """Returns the default config for bring_to_target tasks.""" + config = config_dict.create( + ctrl_dt=0.02, + sim_dt=0.005, + episode_length=150, + action_repeat=1, + action_scale=0.04, + reward_config=config_dict.create( + scales=config_dict.create( + # Gripper goes to the box. + gripper_box=4.0, + # Box goes to the target mocap. + box_target=8.0, + # Do not collide the gripper with the floor. + no_floor_collision=0.25, + # Arm stays close to target pose. + robot_target_qpos=0.3, + ) + ), + impl="jax", + nconmax=24 * 8192, + njmax=128, + ) + return config + class PandaPickCube(panda.PandaBase): - """Bring a box to a target.""" - - def __init__( - self, - config: config_dict.ConfigDict = default_config(), - config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None, - sample_orientation: bool = False, - ): - xml_path = ( - mjx_env.ROOT_PATH - / "manipulation" - / "franka_emika_panda" - / "xmls" - / "mjx_single_cube.xml" - ) - super().__init__( - xml_path, - config, - config_overrides, - ) - self._post_init(obj_name="box", keyframe="home") - self._sample_orientation = sample_orientation - - # Contact sensor IDs. - self._floor_hand_found_sensor = [ - self._mj_model.sensor(f"{geom}_floor_found").id - for geom in ["left_finger_pad", "right_finger_pad", "hand_capsule"] - ] - - def reset(self, rng: jax.Array) -> State: - rng, rng_box, rng_target = jax.random.split(rng, 3) - - # intialize box position - box_pos = ( - jax.random.uniform( - rng_box, - (3,), - minval=jp.array([-0.2, -0.2, 0.0]), - maxval=jp.array([0.2, 0.2, 0.0]), + """Bring a box to a target.""" + + def __init__( + self, + config: config_dict.ConfigDict = default_config(), + config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None, + sample_orientation: bool = False, + ): + xml_path = mjx_env.ROOT_PATH / "manipulation" / "franka_emika_panda" / "xmls" / "mjx_single_cube.xml" + super().__init__( + xml_path, + config, + config_overrides, ) - + self._init_obj_pos - ) + self._post_init(obj_name="box", keyframe="home") + self._sample_orientation = sample_orientation + + # Contact sensor IDs. + self._floor_hand_found_sensor = [ + self._mj_model.sensor(f"{geom}_floor_found").id + for geom in ["left_finger_pad", "right_finger_pad", "hand_capsule"] + ] + + def reset(self, rng: jax.Array) -> State: + rng, rng_box, rng_target = jax.random.split(rng, 3) - # initialize target position - target_pos = ( - jax.random.uniform( - rng_target, - (3,), - minval=jp.array([-0.2, -0.2, 0.2]), - maxval=jp.array([0.2, 0.2, 0.4]), + # intialize box position + box_pos = ( + jax.random.uniform( + rng_box, + (3,), + minval=jp.array([-0.2, -0.2, 0.0]), + maxval=jp.array([0.2, 0.2, 0.0]), + ) + + self._init_obj_pos ) - + self._init_obj_pos - ) - target_quat = jp.array([1.0, 0.0, 0.0, 0.0], dtype=float) - if self._sample_orientation: - # sample a random direction - rng, rng_axis, rng_theta = jax.random.split(rng, 3) - perturb_axis = jax.random.uniform(rng_axis, (3,), minval=-1, maxval=1) - perturb_axis = perturb_axis / math.norm(perturb_axis) - perturb_theta = jax.random.uniform(rng_theta, maxval=np.deg2rad(45)) - target_quat = math.axis_angle_to_quat(perturb_axis, perturb_theta) - - # initialize data - init_q = ( - jp.array(self._init_q) - .at[self._obj_qposadr : self._obj_qposadr + 3] - .set(box_pos) - ) - data = mjx_env.make_data( - self._mj_model, - qpos=init_q, - qvel=jp.zeros(self._mjx_model.nv, dtype=float), - ctrl=self._init_ctrl, - impl=self._mjx_model.impl.value, - nconmax=self._config.nconmax, - njmax=self._config.njmax, - ) + # initialize target position + target_pos = ( + jax.random.uniform( + rng_target, + (3,), + minval=jp.array([-0.2, -0.2, 0.2]), + maxval=jp.array([0.2, 0.2, 0.4]), + ) + + self._init_obj_pos + ) - # set target mocap position - data = data.replace( - mocap_pos=data.mocap_pos.at[self._mocap_target, :].set(target_pos), - mocap_quat=data.mocap_quat.at[self._mocap_target, :].set(target_quat), - ) + target_quat = jp.array([1.0, 0.0, 0.0, 0.0], dtype=float) + if self._sample_orientation: + # sample a random direction + rng, rng_axis, rng_theta = jax.random.split(rng, 3) + perturb_axis = jax.random.uniform(rng_axis, (3,), minval=-1, maxval=1) + perturb_axis = perturb_axis / math.norm(perturb_axis) + perturb_theta = jax.random.uniform(rng_theta, maxval=np.deg2rad(45)) + target_quat = math.axis_angle_to_quat(perturb_axis, perturb_theta) + + # initialize data + init_q = jp.array(self._init_q).at[self._obj_qposadr : self._obj_qposadr + 3].set(box_pos) + data = mjx_env.make_data( + self._mj_model, + qpos=init_q, + qvel=jp.zeros(self._mjx_model.nv, dtype=float), + ctrl=self._init_ctrl, + impl=self._mjx_model.impl.value, + nconmax=self._config.nconmax, + njmax=self._config.njmax, + ) - # initialize env state and info - metrics = { - "out_of_bounds": jp.array(0.0, dtype=float), - **{k: 0.0 for k in self._config.reward_config.scales.keys()}, - } - info = {"rng": rng, "target_pos": target_pos, "reached_box": 0.0} - obs = self._get_obs(data, info) - reward, done = jp.zeros(2) - state = State(data, obs, reward, done, metrics, info) - return state - - def step(self, state: State, action: jax.Array) -> State: - delta = action * self._action_scale - ctrl = state.data.ctrl + delta - ctrl = jp.clip(ctrl, self._lowers, self._uppers) - - data = mjx_env.step(self._mjx_model, state.data, ctrl, self.n_substeps) - - raw_rewards = self._get_reward(data, state.info) - rewards = { - k: v * self._config.reward_config.scales[k] - for k, v in raw_rewards.items() - } - reward = jp.clip(sum(rewards.values()), -1e4, 1e4) - box_pos = data.xpos[self._obj_body] - out_of_bounds = jp.any(jp.abs(box_pos) > 1.0) - out_of_bounds |= box_pos[2] < 0.0 - done = out_of_bounds | jp.isnan(data.qpos).any() | jp.isnan(data.qvel).any() - done = done.astype(float) - - state.metrics.update( - **raw_rewards, out_of_bounds=out_of_bounds.astype(float) - ) + # set target mocap position + data = data.replace( + mocap_pos=data.mocap_pos.at[self._mocap_target, :].set(target_pos), + mocap_quat=data.mocap_quat.at[self._mocap_target, :].set(target_quat), + ) - obs = self._get_obs(data, state.info) - state = State(data, obs, reward, done, state.metrics, state.info) - - return state - - def _get_reward(self, data: mjx.Data, info: Dict[str, Any]) -> Dict[str, Any]: - target_pos = info["target_pos"] - box_pos = data.xpos[self._obj_body] - gripper_pos = data.site_xpos[self._gripper_site] - pos_err = jp.linalg.norm(target_pos - box_pos) - box_mat = data.xmat[self._obj_body] - target_mat = math.quat_to_mat(data.mocap_quat[self._mocap_target]) - rot_err = jp.linalg.norm(target_mat.ravel()[:6] - box_mat.ravel()[:6]) - - box_target = 1 - jp.tanh(5 * (0.9 * pos_err + 0.1 * rot_err)) - gripper_box = 1 - jp.tanh(5 * jp.linalg.norm(box_pos - gripper_pos)) - robot_target_qpos = 1 - jp.tanh( - jp.linalg.norm( - data.qpos[self._robot_arm_qposadr] - - self._init_q[self._robot_arm_qposadr] + # initialize env state and info + metrics = { + "out_of_bounds": jp.array(0.0, dtype=float), + **{k: 0.0 for k in self._config.reward_config.scales.keys()}, + } + info = {"rng": rng, "target_pos": target_pos, "reached_box": 0.0} + obs = self._get_obs(data, info) + reward, done = jp.zeros(2) + state = State(data, obs, reward, done, metrics, info) + return state + + def step(self, state: State, action: jax.Array) -> State: + delta = action * self._action_scale + ctrl = state.data.ctrl + delta + ctrl = jp.clip(ctrl, self._lowers, self._uppers) + + data = mjx_env.step(self._mjx_model, state.data, ctrl, self.n_substeps) + + raw_rewards = self._get_reward(data, state.info) + rewards = {k: v * self._config.reward_config.scales[k] for k, v in raw_rewards.items()} + reward = jp.clip(sum(rewards.values()), -1e4, 1e4) + box_pos = data.xpos[self._obj_body] + out_of_bounds = jp.any(jp.abs(box_pos) > 1.0) + out_of_bounds |= box_pos[2] < 0.0 + done = out_of_bounds | jp.isnan(data.qpos).any() | jp.isnan(data.qvel).any() + done = done.astype(float) + + state.metrics.update(**raw_rewards, out_of_bounds=out_of_bounds.astype(float)) + + obs = self._get_obs(data, state.info) + state = State(data, obs, reward, done, state.metrics, state.info) + + return state + + def _get_reward(self, data: mjx.Data, info: Dict[str, Any]) -> Dict[str, Any]: + target_pos = info["target_pos"] + box_pos = data.xpos[self._obj_body] + gripper_pos = data.site_xpos[self._gripper_site] + pos_err = jp.linalg.norm(target_pos - box_pos) + box_mat = data.xmat[self._obj_body] + target_mat = math.quat_to_mat(data.mocap_quat[self._mocap_target]) + rot_err = jp.linalg.norm(target_mat.ravel()[:6] - box_mat.ravel()[:6]) + + box_target = 1 - jp.tanh(5 * (0.9 * pos_err + 0.1 * rot_err)) + gripper_box = 1 - jp.tanh(5 * jp.linalg.norm(box_pos - gripper_pos)) + robot_target_qpos = 1 - jp.tanh( + jp.linalg.norm(data.qpos[self._robot_arm_qposadr] - self._init_q[self._robot_arm_qposadr]) ) - ) - # Check for collisions with the floor - hand_floor_collision = [ - data.sensordata[self._mj_model.sensor_adr[sensor_id]] > 0 - for sensor_id in self._floor_hand_found_sensor - ] - floor_collision = sum(hand_floor_collision) > 0 - no_floor_collision = (1 - floor_collision).astype(float) - - info["reached_box"] = 1.0 * jp.maximum( - info["reached_box"], - (jp.linalg.norm(box_pos - gripper_pos) < 0.012), - ) + # Check for collisions with the floor + hand_floor_collision = [ + data.sensordata[self._mj_model.sensor_adr[sensor_id]] > 0 for sensor_id in self._floor_hand_found_sensor + ] + floor_collision = sum(hand_floor_collision) > 0 + no_floor_collision = (1 - floor_collision).astype(float) + + info["reached_box"] = 1.0 * jp.maximum( + info["reached_box"], + (jp.linalg.norm(box_pos - gripper_pos) < 0.012), + ) - rewards = { - "gripper_box": gripper_box, - "box_target": box_target * info["reached_box"], - "no_floor_collision": no_floor_collision, - "robot_target_qpos": robot_target_qpos, - } - return rewards - - def _get_obs(self, data: mjx.Data, info: dict[str, Any]) -> jax.Array: - gripper_pos = data.site_xpos[self._gripper_site] - gripper_mat = data.site_xmat[self._gripper_site].ravel() - target_mat = math.quat_to_mat(data.mocap_quat[self._mocap_target]) - obs = jp.concatenate([ - data.qpos, - data.qvel, - gripper_pos, - gripper_mat[3:], - data.xmat[self._obj_body].ravel()[3:], - data.xpos[self._obj_body] - data.site_xpos[self._gripper_site], - info["target_pos"] - data.xpos[self._obj_body], - target_mat.ravel()[:6] - data.xmat[self._obj_body].ravel()[:6], - data.ctrl - data.qpos[self._robot_qposadr[:-1]], - ]) - - return obs + rewards = { + "gripper_box": gripper_box, + "box_target": box_target * info["reached_box"], + "no_floor_collision": no_floor_collision, + "robot_target_qpos": robot_target_qpos, + } + return rewards + + def _get_obs(self, data: mjx.Data, info: dict[str, Any]) -> jax.Array: + gripper_pos = data.site_xpos[self._gripper_site] + gripper_mat = data.site_xmat[self._gripper_site].ravel() + target_mat = math.quat_to_mat(data.mocap_quat[self._mocap_target]) + obs = jp.concatenate([ + data.qpos, + data.qvel, + gripper_pos, + gripper_mat[3:], + data.xmat[self._obj_body].ravel()[3:], + data.xpos[self._obj_body] - data.site_xpos[self._gripper_site], + info["target_pos"] - data.xpos[self._obj_body], + target_mat.ravel()[:6] - data.xmat[self._obj_body].ravel()[:6], + data.ctrl - data.qpos[self._robot_qposadr[:-1]], + ]) + + return obs ``` diff --git a/tools/mjlab_integration/full_runner.py b/tools/mjlab_integration/full_runner.py index 102a70615..18a2eeaec 100644 --- a/tools/mjlab_integration/full_runner.py +++ b/tools/mjlab_integration/full_runner.py @@ -258,7 +258,7 @@ def _align_by_jointname(a: RolloutResult, b: RolloutResult): import mujoco def _short(n: str) -> str: - return n.split("/")[-1] + return n.rsplit("/", maxsplit=1)[-1] a_short = [_short(n) for n in a.joint_names] b_short = [_short(n) for n in b.joint_names]