Adds body acceleration and heading data to RigidObjectData (#200)

# Description

This MR adds rigid body acceleration and heading data into the
`RigidObjectData` class.

## Type of change

- New feature (non-breaking change which adds functionality)
- This change requires a documentation update

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file

---------

Signed-off-by: Mayank Mittal <[email protected]>
Co-authored-by: Nikita Rudin <[email protected]>

source/extensions/omni.isaac.orbit/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.9.14"
+version = "0.9.15"
 
 # Description
 title = "ORBIT framework for Robot Learning"

source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst

@@ -1,6 +1,24 @@
 Changelog
 ---------
 
+0.9.15 (2023-10-22)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added base yaw heading and body acceleration into :class:`omni.isaac.orbit.assets.RigidObjectData` class.
+  These quantities are computed inside the :class:`RigidObject` class.
+
+Fixed
+^^^^^
+
+* Fixed the :meth:`omni.isaac.orbit.assets.RigidObject.set_external_force_and_torque` method to correctly
+  deal with the body indices.
+* Fixed a bug in the :meth:`omni.isaac.orbit.utils.math.wrap_to_pi` method to prevent self-assignment of
+  the input tensor.
+
+
 0.9.14 (2023-10-21)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/rigid_object/rigid_object.py

@@ -104,6 +104,8 @@ def reset(self, env_ids: Sequence[int] | None = None):
         # reset external wrench
         self._external_force_b[env_ids] = 0.0
         self._external_torque_b[env_ids] = 0.0
+        # reset last body vel
+        self._last_body_vel_w[env_ids] = 0.0
 
     def write_data_to_sim(self):
         """Write external wrench to the simulation.
@@ -115,8 +117,8 @@ def write_data_to_sim(self):
         # write external wrench
         if self.has_external_wrench:
             self._body_view._physics_view.apply_forces_and_torques_at_position(
-                force_data=self._external_force_b,
-                torque_data=self._external_torque_b,
+                force_data=self._external_force_b.view(-1, 3),
+                torque_data=self._external_torque_b.view(-1, 3),
                 position_data=None,
                 indices=self._ALL_BODY_INDICES,
                 is_global=False,
@@ -233,7 +235,7 @@ def set_external_force_and_torque(
             env_ids: Environment indices to apply external wrench to. Defaults to None (all instances).
         """
         if forces.any() or torques.any():
-            self.has_external_force = True
+            self.has_external_wrench = True
             # resolve all indices
             # -- env_ids
             if env_ids is None:
@@ -245,16 +247,19 @@ def set_external_force_and_torque(
                 body_ids = torch.arange(self.num_bodies, dtype=torch.long, device=self.device)
             elif not isinstance(body_ids, torch.Tensor):
                 body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device)
+
+            # note: we need to do this complicated indexing since torch doesn't support multi-indexing
             # create global body indices from env_ids and env_body_ids
             # (env_id * total_bodies_per_env) + body_id
             total_bodies_per_env = self.body_view.count // self.root_view.count
             indices = body_ids.repeat(len(env_ids), 1) + env_ids.unsqueeze(1) * total_bodies_per_env
+            indices = indices.view(-1)
             # set into internal buffers
             # note: these are applied in the write_to_sim function
-            self._external_force_b.flatten(0, 1)[indices] = forces
-            self._external_torque_b.flatten(0, 1)[indices] = torques
+            self._external_force_b.flatten(0, 1)[indices] = forces.flatten(0, 1)
+            self._external_torque_b.flatten(0, 1)[indices] = torques.flatten(0, 1)
         else:
-            self.has_external_force = False
+            self.has_external_wrench = False
 
     """
     Internal helper.
@@ -279,6 +284,7 @@ def _create_buffers(self):
         self._ALL_INDICES = torch.arange(self.root_view.count, dtype=torch.long, device=self.device)
         self._ALL_BODY_INDICES = torch.arange(self.body_view.count, dtype=torch.long, device=self.device)
         self._GRAVITY_VEC_W = torch.tensor((0.0, 0.0, -1.0), device=self.device).repeat(self.root_view.count, 1)
+        self._FORWARD_VEC_B = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self.root_view.count, 1)
         # external forces and torques
         self.has_external_wrench = False
         self._external_force_b = torch.zeros((self.root_view.count, self.num_bodies, 3), device=self.device)
@@ -288,13 +294,19 @@ def _create_buffers(self):
         # -- properties
         self._data.body_names = self.body_names
         # -- root states
-        self._data.root_state_w = torch.zeros(self.root_view.count, 13, dtype=torch.float, device=self.device)
+        self._data.root_state_w = torch.zeros(self.root_view.count, 13, device=self.device)
         self._data.default_root_state_w = torch.zeros_like(self._data.root_state_w)
         # -- body states
         self._data.body_state_w = torch.zeros(self.root_view.count, self.num_bodies, 13, device=self.device)
         # -- post-computed
         self._data.root_vel_b = torch.zeros(self.root_view.count, 6, device=self.device)
         self._data.projected_gravity_b = torch.zeros(self.root_view.count, 3, device=self.device)
+        self._data.heading_w = torch.zeros(self.root_view.count, device=self.device)
+        self._data.body_acc_w = torch.zeros(self.root_view.count, self.num_bodies, 6, device=self.device)
+
+        # history buffers for quantities
+        # -- used to compute body accelerations numerically
+        self._last_body_vel_w = torch.zeros(self.root_view.count, self.num_bodies, 6, device=self.device)
 
     def _process_cfg(self):
         """Post processing of configuration parameters."""
@@ -321,6 +333,9 @@ def _update_common_data(self, dt: float):
         self._data.body_state_w[..., :7] = self.body_physx_view.get_transforms().view(-1, self.num_bodies, 7)
         self._data.body_state_w[..., 3:7] = math_utils.convert_quat(self._data.body_state_w[..., 3:7], to="wxyz")
         self._data.body_state_w[..., 7:] = self.body_physx_view.get_velocities().view(-1, self.num_bodies, 6)
+        # -- body acceleration
+        self._data.body_acc_w[:] = (self._data.body_state_w[..., 7:] - self._last_body_vel_w) / dt
+        self._last_body_vel_w[:] = self._data.body_state_w[..., 7:]
         # -- root state in body frame
         self._data.root_vel_b[:, 0:3] = math_utils.quat_rotate_inverse(
             self._data.root_quat_w, self._data.root_lin_vel_w
@@ -329,3 +344,6 @@ def _update_common_data(self, dt: float):
             self._data.root_quat_w, self._data.root_ang_vel_w
         )
         self._data.projected_gravity_b[:] = math_utils.quat_rotate_inverse(self._data.root_quat_w, self._GRAVITY_VEC_W)
+        # -- heading direction of root
+        forward_w = math_utils.quat_apply(self._data.root_quat_w, self._FORWARD_VEC_B)
+        self._data.heading_w[:] = torch.atan2(forward_w[:, 1], forward_w[:, 0])

source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/rigid_object/rigid_object_data.py

@@ -40,10 +40,25 @@ class RigidObjectData:
     projected_gravity_b: torch.Tensor = None
     """Projection of the gravity direction on base frame. Shape is ``(count, 3)``."""
 
+    heading_w: torch.Tensor = None
+    """Yaw heading of the base frame (in radians). Shape is ``(count,)``.
+
+    Note:
+        This quantity is computed by assuming that the forward-direction of the base
+        frame is along x-direction, i.e. :math:`(1, 0, 0)`.
+    """
+
     body_state_w: torch.Tensor = None
     """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame.
     Shape is ``(count, num_bodies, 13)``."""
 
+    body_acc_w: torch.Tensor = None
+    """Acceleration of all bodies. Shape is ``(count, num_bodies, 6)``.
+
+    Note:
+        This quantity is computed based on the rigid body state from the last step.
+    """
+
     """
     Properties
     """
@@ -107,3 +122,13 @@ def body_lin_vel_w(self) -> torch.Tensor:
     def body_ang_vel_w(self) -> torch.Tensor:
         """Angular velocity of all bodies in simulation world frame. Shape is ``(count, num_bodies, 3)``."""
         return self.body_state_w[..., 10:13]
+
+    @property
+    def body_lin_acc_w(self) -> torch.Tensor:
+        """Linear acceleration of all bodies in simulation world frame. Shape is ``(count, num_bodies, 3)``."""
+        return self.body_acc_w[..., 0:3]
+
+    @property
+    def body_ang_acc_w(self) -> torch.Tensor:
+        """Angular acceleration of all bodies in simulation world frame. Shape is ``(count, num_bodies, 3)``."""
+        return self.body_acc_w[..., 3:6]

source/extensions/omni.isaac.orbit/omni/isaac/orbit/command_generators/velocity_command_generator.py

@@ -54,8 +54,6 @@ def __init__(self, cfg: UniformVelocityCommandGeneratorCfg, env: BaseEnv):
         super().__init__(cfg, env)
         # -- robot
         self.robot: Articulation = env.scene[cfg.asset_name]
-        # -- constants
-        self._FORWARD_VEC_B = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self.num_envs, 1)
         # -- command: x vel, y vel, yaw vel, heading
         self.vel_command_b = torch.zeros(self.num_envs, 3, device=self.device)
         self.heading_target = torch.zeros(self.num_envs, device=self.device)
@@ -131,15 +129,10 @@ def _update_command(self):
         # Compute angular velocity from heading direction
         if self.cfg.heading_command:
             # resolve indices of heading envs
-            heading_env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten()
-            # obtain heading direction
-            forward = math_utils.quat_apply(
-                self.robot.data.root_quat_w[heading_env_ids, :], self._FORWARD_VEC_B[heading_env_ids]
-            )
-            heading = torch.atan2(forward[:, 1], forward[:, 0])
+            env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten()
             # compute angular velocity
-            self.vel_command_b[heading_env_ids, 2] = torch.clip(
-                0.5 * math_utils.wrap_to_pi(self.heading_target[heading_env_ids] - heading),
+            self.vel_command_b[env_ids, 2] = torch.clip(
+                0.5 * math_utils.wrap_to_pi(self.heading_target[env_ids] - self.robot.data.heading_w[env_ids]),
                 min=self.cfg.ranges.ang_vel_z[0],
                 max=self.cfg.ranges.ang_vel_z[1],
             )

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/randomizations.py

@@ -123,11 +123,14 @@ def apply_external_force_torque(
     # extract the used quantities (to enable type-hinting)
     asset: RigidObject = env.scene[asset_cfg.name]
     num_envs = env.scene.num_envs
+    # resolve environment ids
+    if env_ids is None:
+        env_ids = torch.arange(num_envs)
 
     # sample random forces and torques
-    body_count = num_envs * len(asset_cfg.body_ids)
-    forces = sample_uniform(*force_range, (body_count, 3), asset.device)
-    torques = sample_uniform(*torque_range, (body_count, 3), asset.device)
+    size = (len(env_ids), len(asset_cfg.body_ids), 3)
+    forces = sample_uniform(*force_range, size, asset.device)
+    torques = sample_uniform(*torque_range, size, asset.device)
     # set the forces and torques into the buffers
     # note: these are only applied when you call: `asset.write_data_to_sim()`
     asset.set_external_force_and_torque(forces, torques, env_ids=env_ids, body_ids=asset_cfg.body_ids)

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/spawners/from_files/from_files.py

@@ -228,7 +228,6 @@ def spawn_ground_plane(prim_path: str, cfg: from_files_cfg.GroundPlaneCfg, **kwa
         else:
             prop_path = f"{prim_path}/Looks/theGrid/Shader.inputs:diffuse_tint"
         # change the color
-        print("Changing color to: ", cfg.color, prop_path)
         omni.kit.commands.execute(
             "ChangePropertyCommand",
             prop_path=Sdf.Path(prop_path),

source/extensions/omni.isaac.orbit/omni/isaac/orbit/utils/math.py

@@ -84,6 +84,7 @@ def wrap_to_pi(angles: torch.Tensor) -> torch.Tensor:
     Returns:
         Angles in the range [-pi, pi].
     """
+    angles = angles.clone()
     angles %= 2 * np.pi
     angles -= 2 * np.pi * (angles > np.pi)
     return angles

source/extensions/omni.isaac.orbit/omni/isaac/orbit/utils/string.py

@@ -215,6 +215,7 @@ def resolve_matching_names(keys: str | Sequence[str], list_of_strings: Sequence[
         msg = "\n"
         for key, value in zip(keys, keys_match_found):
             msg += f"\t{key}: {value}\n"
+        msg += f"Available strings: {list_of_strings}\n"
         # raise error
         raise ValueError(
             f"Not all regular expressions are matched! Please check that the regular expressions are correct: {msg}"
@@ -282,6 +283,7 @@ def resolve_matching_names_values(
         msg = "\n"
         for key, value in zip(data.keys(), keys_match_found):
             msg += f"\t{key}: {value}\n"
+        msg += f"Available strings: {list_of_strings}\n"
         # raise error
         raise ValueError(
             f"Not all regular expressions are matched! Please check that the regular expressions are correct: {msg}"

source/extensions/omni.isaac.orbit/test/assets/check_external_force.py

@@ -0,0 +1,136 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+This script checks if the external force is applied correctly on the robot.
+
+.. code-block:: bash
+
+    # Usage to apply force on base
+    ./orbit.sh -p source/extensions/omni.isaac.orbit/test/assets/check_external_force.py --body base --force 1000
+    # Usage to apply force on legs
+    ./orbit.sh -p source/extensions/omni.isaac.orbit/test/assets/check_external_force.py --body .*_SHANK --force 100
+"""
+
+from __future__ import annotations
+
+"""Launch Isaac Sim Simulator first."""
+
+
+import argparse
+
+from omni.isaac.orbit.app import AppLauncher
+
+# add argparse arguments
+parser = argparse.ArgumentParser(description="This script demonstrates how to external force on a legged robot.")
+parser.add_argument("--body", type=str, help="Name of the body to apply force on.")
+parser.add_argument("--force", type=float, help="Force to apply on the body.")
+# append AppLauncher cli args
+AppLauncher.add_app_launcher_args(parser)
+# parse the arguments
+args_cli = parser.parse_args()
+
+# launch omniverse app
+app_launcher = AppLauncher(args_cli)
+simulation_app = app_launcher.app
+
+"""Rest everything follows."""
+
+import torch
+import traceback
+
+import carb
+
+import omni.isaac.orbit.sim as sim_utils
+from omni.isaac.orbit.assets import Articulation
+from omni.isaac.orbit.assets.config.anymal import ANYMAL_C_CFG
+from omni.isaac.orbit.sim import SimulationContext
+
+
+def main():
+    """Main function."""
+
+    # Load kit helper
+    sim = SimulationContext(sim_utils.SimulationCfg(dt=0.005))
+    # Set main camera
+    sim.set_camera_view(eye=[3.5, 3.5, 3.5], target=[0.0, 0.0, 0.0])
+
+    # Spawn things into stage
+    # Ground-plane
+    cfg = sim_utils.GroundPlaneCfg()
+    cfg.func("/World/defaultGroundPlane", cfg)
+    # Lights
+    cfg = sim_utils.DistantLightCfg(intensity=1000.0, color=(0.75, 0.75, 0.75))
+    cfg.func("/World/Light/greyLight", cfg)
+
+    # Robots
+    robot_cfg = ANYMAL_C_CFG
+    robot_cfg.spawn.func("/World/Anymal_c/Robot_1", robot_cfg.spawn, translation=(0.0, -0.5, 0.65))
+    robot_cfg.spawn.func("/World/Anymal_c/Robot_2", robot_cfg.spawn, translation=(0.0, 0.5, 0.65))
+    # create handles for the robots
+    robot = Articulation(robot_cfg.replace(prim_path="/World/Anymal_c/Robot.*"))
+
+    # Play the simulator
+    sim.reset()
+
+    # Find bodies to apply the force
+    body_ids, body_names = robot.find_bodies(args_cli.body)
+    # Sample a large force
+    external_wrench_b = torch.zeros(robot.root_view.count, len(body_ids), 6, device=sim.device)
+    external_wrench_b[..., 1] = args_cli.force
+
+    # Now we are ready!
+    print("[INFO]: Setup complete...")
+    print("[INFO]: Applying force on the robot: ", args_cli.body, " -> ", body_names)
+
+    # Define simulation stepping
+    sim_dt = sim.get_physics_dt()
+    sim_time = 0.0
+    count = 0
+    # Simulate physics
+    while simulation_app.is_running():
+        # reset
+        if count % 100 == 0:
+            # reset counters
+            sim_time = 0.0
+            count = 0
+            # reset root state
+            root_state = robot.data.default_root_state_w.clone()
+            root_state[0, :2] = torch.tensor([0.0, -0.5], device=sim.device)
+            root_state[1, :2] = torch.tensor([0.0, 0.5], device=sim.device)
+            robot.write_root_state_to_sim(root_state)
+            # reset dof state
+            joint_pos, joint_vel = robot.data.default_joint_pos, robot.data.default_joint_vel
+            robot.write_joint_state_to_sim(joint_pos, joint_vel)
+            robot.reset()
+            # apply force
+            robot.set_external_force_and_torque(
+                external_wrench_b[..., :3], external_wrench_b[..., 3:], body_ids=body_ids
+            )
+            # reset command
+            print(">>>>>>>> Reset!")
+        # apply action to the robot
+        robot.set_joint_position_target(robot.data.default_joint_pos.clone())
+        robot.write_data_to_sim()
+        # perform step
+        sim.step(render=app_launcher.RENDER)
+        # update sim-time
+        sim_time += sim_dt
+        count += 1
+        # update buffers
+        robot.update(sim_dt)
+
+
+if __name__ == "__main__":
+    try:
+        # run the main execution
+        main()
+    except Exception as err:
+        carb.log_error(err)
+        carb.log_error(traceback.format_exc())
+        raise
+    finally:
+        # close sim app
+        simulation_app.close()