Yizew/fix demo scripts and support newton

CONTRIBUTORS.md

@@ -193,6 +193,7 @@ Guidelines for modifications:
 * Yang Jin
 * Yanzi Zhu
 * Yijie Guo
+* Yize Wang
 * Yohan Choi
 * Yujian Zhang
 * Yun Liu

scripts/demos/arl_robot_1.py

@@ -6,22 +6,26 @@
 """
 Script to view ARL Robot 1.
 
-Launch Isaac Sim Simulator first.
+.. code-block:: bash
+
+    # Usage
+    ./isaaclab.sh -p scripts/demos/arl_robot_1.py
+
+    # Run with Newton MJWarp physics and the Newton visualizer
+    ./isaaclab.sh -p scripts/demos/arl_robot_1.py physics=newton_mjwarp --visualizer newton
+
 """
 
 # Create argparser
 import argparse
 
-from isaaclab.app import AppLauncher
+from isaaclab_tasks.utils.demo_launcher import DemoAppLauncher
 
 parser = argparse.ArgumentParser(description="View ARL Robot 1 with Lee Position Controller.")
-# append AppLauncher cli args
-AppLauncher.add_app_launcher_args(parser)
-args_cli = parser.parse_args()
+args_cli = DemoAppLauncher.parse_args(parser)
 
 # launch omniverse app
-app_launcher = AppLauncher(args_cli)
-simulation_app = app_launcher.app
+simulation_app = DemoAppLauncher(args_cli, kit_required=True)
 
 """Rest everything follows."""
 
@@ -45,7 +49,7 @@ def main():
 
     # Create simulation context
     sim_cfg = sim_utils.SimulationCfg(dt=0.01)
-    sim = SimulationContext(sim_cfg)
+    sim = simulation_app.create_context(sim_cfg, SimulationContext)
 
     # Create a dome light with light blue color
     stage = omni.usd.get_context().get_stage()

scripts/demos/arms.py

@@ -11,26 +11,23 @@
     # Usage
     ./isaaclab.sh -p scripts/demos/arms.py
 
-"""
+    # Run with Newton MJWarp physics and the Newton visualizer
+    ./isaaclab.sh -p scripts/demos/arms.py physics=newton_mjwarp --visualizer newton
 
-"""Launch Isaac Sim Simulator first."""
+"""
 
 import argparse
 
-from isaaclab.app import AppLauncher
+from isaaclab_tasks.utils.demo_launcher import DemoAppLauncher
 
 # add argparse arguments
 parser = argparse.ArgumentParser(description="This script demonstrates different single-arm manipulators.")
-# append AppLauncher cli args
-AppLauncher.add_app_launcher_args(parser)
 # demos should open Kit visualizer by default
 parser.set_defaults(visualizer=["kit"])
 # parse the arguments
-args_cli = parser.parse_args()
+args_cli = DemoAppLauncher.parse_args(parser)
+simulation_app = DemoAppLauncher(args_cli)
 
-# launch omniverse app
-app_launcher = AppLauncher(args_cli)
-simulation_app = app_launcher.app
 
 """Rest everything follows."""
 
@@ -91,7 +88,7 @@ def design_scene() -> tuple[dict, list[list[float]]]:
     cfg = sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd")
     cfg.func("/World/Origin1/Table", cfg, translation=(0.55, 0.0, 1.05))
     # -- Robot
-    franka_arm_cfg = FRANKA_PANDA_CFG.replace(prim_path="/World/Origin1/Robot")
+    franka_arm_cfg = simulation_app.tune_articulation_cfg(FRANKA_PANDA_CFG.replace(prim_path="/World/Origin1/Robot"))
     franka_arm_cfg.init_state.pos = (0.0, 0.0, 1.05)
     franka_panda = Articulation(cfg=franka_arm_cfg)
 
@@ -103,7 +100,7 @@ def design_scene() -> tuple[dict, list[list[float]]]:
     )
     cfg.func("/World/Origin2/Table", cfg, translation=(0.0, 0.0, 1.03))
     # -- Robot
-    ur10_cfg = UR10_CFG.replace(prim_path="/World/Origin2/Robot")
+    ur10_cfg = simulation_app.tune_articulation_cfg(UR10_CFG.replace(prim_path="/World/Origin2/Robot"))
     ur10_cfg.init_state.pos = (0.0, 0.0, 1.03)
     ur10 = Articulation(cfg=ur10_cfg)
 
@@ -113,7 +110,9 @@ def design_scene() -> tuple[dict, list[list[float]]]:
     cfg = sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/ThorlabsTable/table_instanceable.usd")
     cfg.func("/World/Origin3/Table", cfg, translation=(0.0, 0.0, 0.8))
     # -- Robot
-    kinova_arm_cfg = KINOVA_JACO2_N7S300_CFG.replace(prim_path="/World/Origin3/Robot")
+    kinova_arm_cfg = simulation_app.tune_articulation_cfg(
+        KINOVA_JACO2_N7S300_CFG.replace(prim_path="/World/Origin3/Robot")
+    )
     kinova_arm_cfg.init_state.pos = (0.0, 0.0, 0.8)
     kinova_j2n7s300 = Articulation(cfg=kinova_arm_cfg)
 
@@ -123,7 +122,9 @@ def design_scene() -> tuple[dict, list[list[float]]]:
     cfg = sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/ThorlabsTable/table_instanceable.usd")
     cfg.func("/World/Origin4/Table", cfg, translation=(0.0, 0.0, 0.8))
     # -- Robot
-    kinova_arm_cfg = KINOVA_JACO2_N6S300_CFG.replace(prim_path="/World/Origin4/Robot")
+    kinova_arm_cfg = simulation_app.tune_articulation_cfg(
+        KINOVA_JACO2_N6S300_CFG.replace(prim_path="/World/Origin4/Robot")
+    )
     kinova_arm_cfg.init_state.pos = (0.0, 0.0, 0.8)
     kinova_j2n6s300 = Articulation(cfg=kinova_arm_cfg)
 
@@ -133,7 +134,7 @@ def design_scene() -> tuple[dict, list[list[float]]]:
     cfg = sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd")
     cfg.func("/World/Origin5/Table", cfg, translation=(0.55, 0.0, 1.05))
     # -- Robot
-    kinova_arm_cfg = KINOVA_GEN3_N7_CFG.replace(prim_path="/World/Origin5/Robot")
+    kinova_arm_cfg = simulation_app.tune_articulation_cfg(KINOVA_GEN3_N7_CFG.replace(prim_path="/World/Origin5/Robot"))
     kinova_arm_cfg.init_state.pos = (0.0, 0.0, 1.05)
     kinova_gen3n7 = Articulation(cfg=kinova_arm_cfg)
 
@@ -145,7 +146,7 @@ def design_scene() -> tuple[dict, list[list[float]]]:
     )
     cfg.func("/World/Origin6/Table", cfg, translation=(0.0, 0.0, 1.03))
     # -- Robot
-    sawyer_arm_cfg = SAWYER_CFG.replace(prim_path="/World/Origin6/Robot")
+    sawyer_arm_cfg = simulation_app.tune_articulation_cfg(SAWYER_CFG.replace(prim_path="/World/Origin6/Robot"))
     sawyer_arm_cfg.init_state.pos = (0.0, 0.0, 1.03)
     sawyer = Articulation(cfg=sawyer_arm_cfg)
 
@@ -216,7 +217,7 @@ def main():
     """Main function."""
     # Initialize the simulation context
     sim_cfg = sim_utils.SimulationCfg(device=args_cli.device)
-    sim = sim_utils.SimulationContext(sim_cfg)
+    sim = simulation_app.create_context(sim_cfg, sim_utils.SimulationContext)
     # Set main camera
     sim.set_camera_view([3.5, 0.0, 3.2], [0.0, 0.0, 0.5])
     # design scene
@@ -231,7 +232,5 @@ def main():
 
 
 if __name__ == "__main__":
-    # run the main function
     main()
-    # close sim app
     simulation_app.close()

scripts/demos/bin_packing.py

@@ -16,6 +16,9 @@
     # Usage
     ./isaaclab.sh -p scripts/demos/bin_packing.py --num_envs 32
 
+    # Run with Newton MJWarp physics and the Newton visualizer
+    ./isaaclab.sh -p scripts/demos/bin_packing.py physics=newton_mjwarp --visualizer newton
+
 """
 
 from __future__ import annotations
@@ -25,27 +28,25 @@
 
 import argparse
 
-from isaaclab.app import AppLauncher
+from isaaclab_tasks.utils.demo_launcher import DemoAppLauncher
 
 # add argparse arguments
 parser = argparse.ArgumentParser(description="Demo usage of RigidObjectCollection through bin packing example")
 parser.add_argument("--num_envs", type=int, default=16, help="Number of environments to spawn.")
-# append AppLauncher cli args
-AppLauncher.add_app_launcher_args(parser)
 # demos should open Kit visualizer by default
 parser.set_defaults(visualizer=["kit"])
 # parse the arguments
-args_cli = parser.parse_args()
+args_cli = DemoAppLauncher.parse_args(parser)
 
 # launch omniverse app
-app_launcher = AppLauncher(args_cli)
-simulation_app = app_launcher.app
+simulation_app = DemoAppLauncher(args_cli)
 
 """Rest everything follows."""
 
 import math
 
 import torch
+import warp as wp
 
 import isaaclab.sim as sim_utils
 import isaaclab.utils.math as math_utils
@@ -141,6 +142,21 @@ class MultiObjectSceneCfg(InteractiveSceneCfg):
     )
 
 
+def torch_indices_to_warp(indices: torch.Tensor, device: str | None = None) -> wp.array:
+    """Convert Torch indices to a Warp int32 array."""
+    indices = indices.to(dtype=torch.int32).contiguous()
+    return wp.clone(wp.from_torch(indices, dtype=wp.int32), device=device or str(indices.device))
+
+
+def reshape_data_to_view_tensor(
+    rigid_object_collection: RigidObjectCollection, data: torch.Tensor, data_dim: int, device: str
+) -> torch.Tensor:
+    """Reshape Torch data from collection order to PhysX view order."""
+    data = data.contiguous()
+    data_wp = wp.from_torch(data, dtype=wp.float32)
+    return wp.to_torch(rigid_object_collection.reshape_data_to_view_3d(data_wp, data_dim, device=device))
+
+
 def reset_object_collections(
     scene: InteractiveScene, asset_name: str, view_states: torch.Tensor, view_ids: torch.Tensor, noise: bool = False
 ) -> None:
@@ -189,8 +205,11 @@ def reset_object_collections(
     view_states[view_ids, :7] = torch.concat((positions, orientations), dim=-1)
     view_states[view_ids, 7:] = new_velocities
 
-    rigid_object_collection.root_view.set_transforms(view_states[:, :7], indices=view_ids)
-    rigid_object_collection.root_view.set_velocities(view_states[:, 7:], indices=view_ids)
+    view_ids_wp = torch_indices_to_warp(view_ids, device=scene.device)
+    transforms_wp = wp.from_torch(view_states[:, :7].contiguous(), dtype=wp.float32)
+    velocities_wp = wp.from_torch(view_states[:, 7:].contiguous(), dtype=wp.float32)
+    rigid_object_collection.root_view.set_transforms(transforms_wp, indices=view_ids_wp)
+    rigid_object_collection.root_view.set_velocities(velocities_wp, indices=view_ids_wp)
 
 
 def build_grocery_defaults(
@@ -280,32 +299,37 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene) -> None:
     # Offset poses into each environment's world frame.
     active_spawn_poses[..., :3] += scene.env_origins.view(-1, 1, 3)
     cached_spawn_poses[..., :3] += scene.env_origins.view(-1, 1, 3)
-    active_spawn_poses = groceries.reshape_data_to_view(active_spawn_poses)
-    cached_spawn_poses = groceries.reshape_data_to_view(cached_spawn_poses)
-    spawn_w = groceries.reshape_data_to_view(default_state_w).clone()
+    active_spawn_poses = reshape_data_to_view_tensor(groceries, active_spawn_poses, 7, device)
+    cached_spawn_poses = reshape_data_to_view_tensor(groceries, cached_spawn_poses, 7, device)
+    spawn_w = reshape_data_to_view_tensor(groceries, default_state_w, 13, device).clone()
 
-    groceries_mask_helper = torch.arange(num_objects * num_envs, device=device) % num_objects
-    # Precompute a helper mask to toggle objects between active and cached sets.
+    # Precompute helpers in PhysX view order: body_0/env_0, body_0/env_1, ..., body_1/env_0, ...
+    view_env_ids = torch.arange(num_envs, device=device).repeat(num_objects)
+    view_object_ids = torch.arange(num_objects, device=device).repeat_interleave(num_envs)
     # Precompute XY bounds [[x_min,y_min],[x_max,y_max]]
     bounds_xy = torch.as_tensor(BIN_XY_BOUND, device=device, dtype=spawn_w.dtype)
+    view_indices_cpu_wp = wp.array(list(range(num_envs * num_objects)), dtype=wp.int32, device="cpu")
     # Simulation loop
     while simulation_app.is_running():
         # Reset
         if count % 250 == 0:
             # reset counter
             count = 0
             # Randomly choose how many groceries stay active in each environment.
-            num_active_groceries = torch.randint(MIN_OBJECTS_PER_BIN, num_objects, (num_envs, 1), device=device)
-            groceries_mask = (groceries_mask_helper.view(num_envs, -1) < num_active_groceries).view(-1, 1)
+            num_active_groceries = torch.randint(MIN_OBJECTS_PER_BIN, num_objects, (num_envs,), device=device)
+            groceries_mask = (view_object_ids < num_active_groceries[view_env_ids]).unsqueeze(-1)
             spawn_w[:, :7] = cached_spawn_poses * (~groceries_mask) + active_spawn_poses * groceries_mask
             # Retrieve positions
             with Timer("[INFO] Time to reset scene: "):
                 reset_object_collections(scene, "groceries", spawn_w, view_indices[~groceries_mask.view(-1)])
                 reset_object_collections(scene, "groceries", spawn_w, view_indices[groceries_mask.view(-1)], noise=True)
                 # Vary the mass and gravity settings so cached objects stay parked.
-                random_masses = torch.rand(groceries.num_instances * num_objects, device=device) * 0.2 + 0.2
-                groceries.root_view.set_masses(random_masses.cpu(), view_indices.cpu())
-                groceries.root_view.set_disable_gravities((~groceries_mask).cpu(), indices=view_indices.cpu())
+                random_masses = torch.rand((groceries.num_instances * num_objects, 1), device="cpu") * 0.2 + 0.2
+                random_masses_wp = wp.from_torch(random_masses.contiguous(), dtype=wp.float32)
+                groceries.root_view.set_masses(random_masses_wp, indices=view_indices_cpu_wp)
+                disable_gravities = (~groceries_mask).cpu().to(torch.uint8).contiguous()
+                disable_gravities_wp = wp.from_torch(disable_gravities, dtype=wp.uint8)
+                groceries.root_view.set_disable_gravities(disable_gravities_wp, indices=view_indices_cpu_wp)
                 scene.reset()
 
         # Write data to sim
@@ -314,7 +338,8 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene) -> None:
         sim.step()
 
         # Bring out-of-bounds objects back to the bin in one pass.
-        xy = groceries.reshape_data_to_view(groceries.data.object_pos_w - scene.env_origins.unsqueeze(1))[:, :2]
+        body_pos_w = groceries.data.body_link_pos_w.torch - scene.env_origins.unsqueeze(1)
+        xy = reshape_data_to_view_tensor(groceries, body_pos_w, 3, device)[:, :2]
         out_bound = torch.nonzero(~((xy >= bounds_xy[0]) & (xy <= bounds_xy[1])).all(dim=1), as_tuple=False).flatten()
         if out_bound.numel():
             # Teleport stray objects back into the active stack to keep the bin tidy.
@@ -333,7 +358,7 @@ def main() -> None:
     """
     # Load kit helper
     sim_cfg = sim_utils.SimulationCfg(dt=0.005, device=args_cli.device)
-    sim = SimulationContext(sim_cfg)
+    sim = simulation_app.create_context(sim_cfg, SimulationContext)
     # Set main camera
     sim.set_camera_view((2.5, 0.0, 4.0), (0.0, 0.0, 2.0))
 

scripts/demos/bipeds.py

@@ -11,26 +11,24 @@
     # Usage
     ./isaaclab.sh -p scripts/demos/bipeds.py
 
-"""
+    # Run with Newton MJWarp physics and the Newton visualizer
+    ./isaaclab.sh -p scripts/demos/bipeds.py physics=newton_mjwarp --visualizer newton
 
-"""Launch Isaac Sim Simulator first."""
+"""
 
 import argparse
 
-from isaaclab.app import AppLauncher
+from isaaclab_tasks.utils.demo_launcher import DemoAppLauncher
 
 # add argparse arguments
 parser = argparse.ArgumentParser(description="This script demonstrates how to simulate bipedal robots.")
-# append AppLauncher cli args
-AppLauncher.add_app_launcher_args(parser)
 # demos should open Kit visualizer by default
 parser.set_defaults(visualizer=["kit"])
 # parse the arguments
-args_cli = parser.parse_args()
+args_cli = DemoAppLauncher.parse_args(parser)
 
 # launch omniverse app
-app_launcher = AppLauncher(args_cli)
-simulation_app = app_launcher.app
+simulation_app = DemoAppLauncher(args_cli)
 
 """Rest everything follows."""
 
@@ -122,7 +120,7 @@ def main():
     """Main function."""
     # Load kit helper
     sim_cfg = sim_utils.SimulationCfg(dt=0.005, device=args_cli.device)
-    sim = SimulationContext(sim_cfg)
+    sim = simulation_app.create_context(sim_cfg, SimulationContext)
     # Set main camera
     sim.set_camera_view(eye=[3.0, 0.0, 2.25], target=[0.0, 0.0, 1.0])
 

scripts/demos/deformables.py

@@ -12,24 +12,20 @@
 
 """
 
-"""Launch Isaac Sim Simulator first."""
-
+# TODO(yizew@nvidia.com): Non-rigid bodies are not supported by the Newton backend yet.
 
 import argparse
 
-from isaaclab.app import AppLauncher
+from isaaclab_tasks.utils.demo_launcher import DemoAppLauncher
 
 # create argparser
 parser = argparse.ArgumentParser(description="This script demonstrates how to spawn deformable prims into the scene.")
-# append AppLauncher cli args
-AppLauncher.add_app_launcher_args(parser)
 # demos should open Kit visualizer by default
 parser.set_defaults(visualizer=["kit"])
 # parse the arguments
-args_cli = parser.parse_args()
+args_cli = DemoAppLauncher.parse_args(parser)
 # launch omniverse app
-app_launcher = AppLauncher(args_cli)
-simulation_app = app_launcher.app
+simulation_app = DemoAppLauncher(args_cli, kit_required=True)
 
 """Rest everything follows."""
 
@@ -225,7 +221,7 @@ def main():
     """Main function."""
     # Initialize the simulation context
     sim_cfg = sim_utils.SimulationCfg(dt=0.01, device=args_cli.device)
-    sim = sim_utils.SimulationContext(sim_cfg)
+    sim = simulation_app.create_context(sim_cfg, sim_utils.SimulationContext)
     # Set main camera
     sim.set_camera_view([4.0, 4.0, 3.0], [0.5, 0.5, 0.0])