object Velocity class like Pose

isaaclab_arena/assets/object_base.py

@@ -17,6 +17,7 @@
 from isaaclab_arena.relations.relations import AtPosition, Relation, RelationBase
 from isaaclab_arena.terms.events import set_object_pose
 from isaaclab_arena.utils.pose import Pose, PoseRange
+from isaaclab_arena.utils.velocity import Velocity
 
 
 class ObjectType(Enum):
@@ -42,7 +43,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.initial_velocity: Velocity | None = None
         self.object_cfg = None
         self.event_cfg = None
         self.relations: list[RelationBase] = []
@@ -81,19 +82,22 @@ 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.
+    def set_initial_velocity(self, velocity: Velocity) -> None:
+        """Set / override the initial 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
+        The velocity is applied as ``init_state.lin_vel`` and
+        ``init_state.ang_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.
+            velocity: A ``Velocity`` specifying linear and angular components.
         """
         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.object_cfg.init_state.lin_vel = velocity.linear_xyz
+        if self.object_cfg is not None and hasattr(self.object_cfg.init_state, "ang_vel"):
+            self.object_cfg.init_state.ang_vel = velocity.angular_xyz
         self.event_cfg = self._init_event_cfg()
 
     def _requires_reset_pose_event(self) -> bool:

isaaclab_arena/terms/events.py

@@ -9,14 +9,15 @@
 from isaaclab.managers import SceneEntityCfg
 
 from isaaclab_arena.utils.pose import Pose
+from isaaclab_arena.utils.velocity import Velocity
 
 
 def set_object_pose(
     env: ManagerBasedEnv,
     env_ids: torch.Tensor,
     asset_cfg: SceneEntityCfg,
     pose: Pose,
-    velocity: tuple[float, float, float] | None = None,
+    velocity: Velocity | None = None,
 ) -> None:
     if env_ids is None:
         return
@@ -29,8 +30,7 @@ def set_object_pose(
     # Set the pose and velocity
     asset.write_root_pose_to_sim(pose_t_xyz_q_wxyz, 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)
+        vel = velocity.to_tensor(device=env.device).unsqueeze(0).expand(num_envs, -1)
         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)

isaaclab_arena/tests/test_events.py

@@ -120,6 +120,7 @@ def _test_object_moves_with_initial_velocity(simulation_app):
     from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment
     from isaaclab_arena.scene.scene import Scene
     from isaaclab_arena.utils.pose import Pose
+    from isaaclab_arena.utils.velocity import Velocity
 
     asset_registry = AssetRegistry()
     no_gravity_cfg = sim_utils.SphereCfg(
@@ -138,9 +139,9 @@ def _test_object_moves_with_initial_velocity(simulation_app):
     )
     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
+    initial_velocity = Velocity(linear_xyz=(-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)
+    sphere.set_initial_velocity(initial_velocity)
 
     scene = Scene(assets=[sphere])
     isaaclab_arena_environment = IsaacLabArenaEnvironment(

isaaclab_arena/utils/velocity.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from dataclasses import dataclass
+
+
+@dataclass
+class Velocity:
+    """Linear and angular velocity of a rigid body in the world frame."""
+
+    linear_xyz: tuple[float, float, float] = (0.0, 0.0, 0.0)
+    """Linear velocity (vx, vy, vz) in the world frame."""
+
+    angular_xyz: tuple[float, float, float] = (0.0, 0.0, 0.0)
+    """Angular velocity (wx, wy, wz) in the world frame."""
+
+    def __post_init__(self):
+        assert isinstance(self.linear_xyz, tuple)
+        assert isinstance(self.angular_xyz, tuple)
+        assert len(self.linear_xyz) == 3
+        assert len(self.angular_xyz) == 3
+
+    @staticmethod
+    def zero() -> "Velocity":
+        return Velocity(linear_xyz=(0.0, 0.0, 0.0), angular_xyz=(0.0, 0.0, 0.0))
+
+    def to_tensor(self, device: torch.device) -> torch.Tensor:
+        """Convert the velocity to a tensor.
+
+        The returned tensor has shape (6,), ordered as (vx, vy, vz, wx, wy, wz).
+
+        Args:
+            device: The device to place the tensor on.
+
+        Returns:
+            The velocity as a tensor of shape (6,).
+        """
+        linear_tensor = torch.tensor(self.linear_xyz, device=device)
+        angular_tensor = torch.tensor(self.angular_xyz, device=device)
+        return torch.cat([linear_tensor, angular_tensor])