Sphere and initial object velocity (#479)

## Summary
New sphere in object library and ObjectBase has now optional initial
velocity

## Detailed description
Now it can be spawn spheres, and objects like the sphere can be
initialized with linear velocity.

Author: DavidRecasens

isaaclab_arena/assets/object_base.py

@@ -42,6 +42,7 @@ def __init__(
         self.prim_path = prim_path
         self.object_type = object_type
         self.initial_pose: Pose | PoseRange | None = None
+        self.initial_velocity: tuple[float, float, float] | None = None
         self.object_cfg = None
         self.event_cfg = None
         self.relations: list[RelationBase] = []
@@ -80,6 +81,21 @@ def set_initial_pose(self, pose: Pose | PoseRange) -> None:
             self.object_cfg.init_state.rot = initial_pose.rotation_wxyz
         self.event_cfg = self._init_event_cfg()
 
+    def set_initial_linear_velocity(self, velocity: tuple[float, float, float]) -> None:
+        """Set / override the initial linear velocity and rebuild derived configs.
+
+        The velocity is applied as the ``init_state.lin_vel`` on the underlying
+        config (``RigidObjectCfg`` or ``ArticulationCfg``) and is also restored
+        on every environment reset via the reset event.
+
+        Args:
+            velocity: Linear velocity ``(vx, vy, vz)`` in the world frame.
+        """
+        self.initial_velocity = velocity
+        if self.object_cfg is not None and hasattr(self.object_cfg.init_state, "lin_vel"):
+            self.object_cfg.init_state.lin_vel = velocity
+        self.event_cfg = self._init_event_cfg()
+
     def _requires_reset_pose_event(self) -> bool:
         """Whether a reset-event for the initial pose should be generated.
 
@@ -91,7 +107,7 @@ def _requires_reset_pose_event(self) -> bool:
         )
 
     def _init_event_cfg(self) -> EventTermCfg | None:
-        """Build the ``EventTermCfg`` for resetting this object's pose."""
+        """Build the ``EventTermCfg`` for resetting this object's pose and velocity."""
         if not self._requires_reset_pose_event():
             return None
 
@@ -106,13 +122,15 @@ def _init_event_cfg(self) -> EventTermCfg | None:
                 },
             )
         else:
+            params: dict = {
+                "pose": initial_pose,
+                "asset_cfg": SceneEntityCfg(self.name),
+                "velocity": self.initial_velocity,
+            }
             return EventTermCfg(
                 func=set_object_pose,
                 mode="reset",
-                params={
-                    "pose": initial_pose,
-                    "asset_cfg": SceneEntityCfg(self.name),
-                },
+                params=params,
             )
 
     def get_relations(self) -> list[RelationBase]:

isaaclab_arena/assets/object_library.py

@@ -10,6 +10,7 @@
 
 if TYPE_CHECKING:
     from isaaclab_arena.assets.hdr_image import HDRImage
+from isaaclab.assets import RigidObjectCfg
 from isaaclab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg
 from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR
 
@@ -419,6 +420,61 @@ def __init__(
         super().__init__(instance_name=instance_name, prim_path=prim_path, initial_pose=initial_pose)
 
 
+@register_asset
+class Sphere(LibraryObject):
+    """
+    A sphere with rigid body physics (dynamic by default).
+    """
+
+    name = "sphere"
+    tags = ["object"]
+    object_type = ObjectType.SPAWNER
+    scale = (1.0, 1.0, 1.0)
+    default_spawner_cfg = sim_utils.SphereCfg(
+        radius=0.1,
+        visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.8, 0.2, 0.2)),
+        collision_props=sim_utils.CollisionPropertiesCfg(),
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+            solver_position_iteration_count=16,
+            solver_velocity_iteration_count=1,
+            max_angular_velocity=1000.0,
+            max_linear_velocity=1000.0,
+            max_depenetration_velocity=5.0,
+            disable_gravity=False,
+        ),
+        mass_props=sim_utils.MassPropertiesCfg(mass=0.25),
+    )
+
+    def __init__(
+        self,
+        instance_name: str | None = None,
+        prim_path: str | None = None,
+        initial_pose: Pose | None = None,
+        scale: tuple[float, float, float] | None = None,
+        spawner_cfg: sim_utils.SphereCfg = default_spawner_cfg,
+    ):
+        self.spawner_cfg = spawner_cfg
+        super().__init__(
+            instance_name=instance_name,
+            prim_path=prim_path,
+            initial_pose=initial_pose,
+            scale=scale,
+        )
+
+    def _generate_spawner_cfg(self) -> RigidObjectCfg:
+        object_cfg = RigidObjectCfg(
+            prim_path=self.prim_path,
+            spawn=self.spawner_cfg,
+        )
+        object_cfg = self._add_initial_pose_to_cfg(object_cfg)
+        if self.initial_velocity is not None:
+            object_cfg.init_state.lin_vel = self.initial_velocity
+        return object_cfg
+
+    def _requires_reset_pose_event(self) -> bool:
+        return self.get_initial_pose() is not None and self.reset_pose
+
+
 @register_asset
 class DomeLight(LibraryObject):
     """A dome light, optionally textured with an HDR image environment map.

isaaclab_arena/terms/events.py

@@ -16,6 +16,7 @@ def set_object_pose(
     env_ids: torch.Tensor,
     asset_cfg: SceneEntityCfg,
     pose: Pose,
+    velocity: tuple[float, float, float] | None = None,
 ) -> None:
     if env_ids is None:
         return
@@ -27,7 +28,12 @@ def set_object_pose(
     pose_t_xyz_q_wxyz[:, :3] += env.scene.env_origins[env_ids]
     # Set the pose and velocity
     asset.write_root_pose_to_sim(pose_t_xyz_q_wxyz, env_ids=env_ids)
-    asset.write_root_velocity_to_sim(torch.zeros(1, 6, device=env.device), env_ids=env_ids)
+    if velocity is not None:
+        vel = torch.zeros(num_envs, 6, device=env.device)
+        vel[:, :3] = torch.tensor(velocity, device=env.device)
+        asset.write_root_velocity_to_sim(vel, env_ids=env_ids)
+    else:
+        asset.write_root_velocity_to_sim(torch.zeros(1, 6, device=env.device), env_ids=env_ids)
 
 
 def set_object_pose_per_env(

isaaclab_arena/tests/test_detect_object_type.py

@@ -72,7 +72,7 @@ def _test_detect_object_type_for_all_objects(simulation_app):
         # the simple RIGID/ARTICULATION classification:
         # - For example, the "peg" and "hole" assets have both RigidBodyAPI and ArticulationRootAPI
         #   applied simultaneously, sometimes in different prim layers.
-        if object_asset.name not in ("hole", "peg", "small_gear", "medium_gear", "large_gear", "gear_base"):
+        if object_asset.name not in ("hole", "peg", "small_gear", "medium_gear", "large_gear", "gear_base", "sphere"):
             print(f"Automatically classifying: {object_asset.name}")
             detected_object_type = detect_object_type(usd_path=object_asset.usd_path)
             print(f"database object type: {object_asset.object_type}")

isaaclab_arena/tests/test_events.py

@@ -12,6 +12,9 @@
 HEADLESS = True
 INITIAL_POSITION_EPS = 0.1  # The cracker box falls slightly.
 
+# For object initial velocity test: min displacement in velocity direction
+OBJECT_VELOCITY_MIN_DISPLACEMENT = 0.05  # minimum movement in meters
+
 
 def _test_set_object_pose_per_env_event(simulation_app):
     """Returns a scene which we use for these tests."""
@@ -107,6 +110,76 @@ def _test_set_object_pose_per_env_event(simulation_app):
     return True
 
 
+def _test_object_moves_with_initial_velocity(simulation_app):
+    """Test that a sphere moves with the given initial velocity after reset."""
+    import isaaclab.sim as sim_utils
+
+    from isaaclab_arena.assets.asset_registry import AssetRegistry
+    from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser
+    from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder
+    from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment
+    from isaaclab_arena.scene.scene import Scene
+    from isaaclab_arena.utils.pose import Pose
+
+    asset_registry = AssetRegistry()
+    no_gravity_cfg = sim_utils.SphereCfg(
+        radius=0.1,
+        visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.8, 0.2, 0.2)),
+        collision_props=sim_utils.CollisionPropertiesCfg(),
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+            solver_position_iteration_count=16,
+            solver_velocity_iteration_count=1,
+            max_angular_velocity=1000.0,
+            max_linear_velocity=1000.0,
+            max_depenetration_velocity=5.0,
+            disable_gravity=True,
+        ),
+        mass_props=sim_utils.MassPropertiesCfg(mass=0.25),
+    )
+    sphere = asset_registry.get_asset_by_name("sphere")(spawner_cfg=no_gravity_cfg)
+
+    initial_velocity = (-0.5, 0.0, 0.0)  # There is a wall in +x
+    sphere.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.5), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)))
+    sphere.set_initial_linear_velocity(initial_velocity)
+
+    scene = Scene(assets=[sphere])
+    isaaclab_arena_environment = IsaacLabArenaEnvironment(
+        name="object_initial_velocity_test",
+        scene=scene,
+    )
+
+    args_cli = get_isaaclab_arena_cli_parser().parse_args([])
+    args_cli.num_envs = 1
+    env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli)
+    env = env_builder.make_registered()
+    env.reset()
+
+    try:
+        initial_position = env.unwrapped.scene[sphere.name].data.root_pose_w[0, :3].clone()
+        initial_position[:3] -= env.unwrapped.scene.env_origins[0]
+
+        for _ in tqdm.tqdm(range(NUM_STEPS)):
+            with torch.inference_mode():
+                actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device)
+                env.step(actions)
+
+        final_position = env.unwrapped.scene[sphere.name].data.root_pose_w[0, :3].clone()
+        final_position[:3] -= env.unwrapped.scene.env_origins[0]
+        displacement = final_position - initial_position
+        assert (
+            displacement[0].item() < -OBJECT_VELOCITY_MIN_DISPLACEMENT  # - x because initial velocity is in -x
+        ), f"Object did not move with given velocity: displacement in x was {displacement[0].item()}"
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+    finally:
+        env.close()
+
+    return True
+
+
 def test_set_object_post_per_env_event():
     result = run_simulation_app_function(
         _test_set_object_pose_per_env_event,
@@ -115,5 +188,14 @@ def test_set_object_post_per_env_event():
     assert result, f"Test {test_set_object_post_per_env_event.__name__} failed"
 
 
+def test_object_moves_with_initial_velocity():
+    result = run_simulation_app_function(
+        _test_object_moves_with_initial_velocity,
+        headless=HEADLESS,
+    )
+    assert result, f"Test {test_object_moves_with_initial_velocity.__name__} failed"
+
+
 if __name__ == "__main__":
     test_set_object_post_per_env_event()
+    test_object_moves_with_initial_velocity()