From b8a2f4648a97dee0db8730c91ceb3aabee068620 Mon Sep 17 00:00:00 2001 From: DavidRecasens Date: Fri, 20 Mar 2026 17:34:08 +0100 Subject: [PATCH] object Velocity class like Pose --- isaaclab_arena/assets/object_base.py | 18 +++++++----- isaaclab_arena/terms/events.py | 6 ++-- isaaclab_arena/tests/test_events.py | 5 ++-- isaaclab_arena/utils/velocity.py | 43 ++++++++++++++++++++++++++++ 4 files changed, 60 insertions(+), 12 deletions(-) create mode 100644 isaaclab_arena/utils/velocity.py diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py index 306d0bfec..85585c6ca 100644 --- a/isaaclab_arena/assets/object_base.py +++ b/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: diff --git a/isaaclab_arena/terms/events.py b/isaaclab_arena/terms/events.py index a62930392..3c830128a 100644 --- a/isaaclab_arena/terms/events.py +++ b/isaaclab_arena/terms/events.py @@ -9,6 +9,7 @@ from isaaclab.managers import SceneEntityCfg from isaaclab_arena.utils.pose import Pose +from isaaclab_arena.utils.velocity import Velocity def set_object_pose( @@ -16,7 +17,7 @@ def set_object_pose( 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) diff --git a/isaaclab_arena/tests/test_events.py b/isaaclab_arena/tests/test_events.py index 8dde818ab..d87bb4488 100644 --- a/isaaclab_arena/tests/test_events.py +++ b/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( diff --git a/isaaclab_arena/utils/velocity.py b/isaaclab_arena/utils/velocity.py new file mode 100644 index 000000000..17f77414c --- /dev/null +++ b/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])