Fixes training of anymal rough terrain locomotion (#207)

# Description

Fixes the environment standalone scripts and rough terrain locomotion
training.


## Type of change

- Bug fix (non-breaking change which fixes an issue)


## 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
- [ ] 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

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

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.9.20"
+version = "0.9.21"
 
 # Description
 title = "ORBIT framework for Robot Learning"

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

 Changelog
 ---------
 
-0.9.20 (2023-10-03)
+0.9.21 (2023-10-26)
+~~~~~~~~~~~~~~~~~~~
+
+Fixed
+^^^^^
+
+* Decreased the priority of callbacks in asset and sensor base classes. This may help in preventing
+  crashes when warm starting the simulation.
+* Fixed no rendering mode when running the environment from the GUI. Earlier the function
+  :meth:`SimulationContext.set_render_mode` was erroring out.
+
+
+0.9.20 (2023-10-25)
 ~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -24,8 +36,8 @@ Added
 * Added Gym observation and action spaces for the :class:`omni.isaac.orbit.envs.RLEnv` class.
 
 
-0.9.18 (2023-10-19)
-~~~~~~~~~~~~~~~~~~
+0.9.18 (2023-10-23)
+~~~~~~~~~~~~~~~~~~~
 
 Added
 ^^^^^

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

@@ -56,16 +56,19 @@ class AssetBase(ABC):
 
         # note: Use weakref on all callbacks to ensure that this object can be deleted when its destructor is called.
         # add callbacks for stage play/stop
+        # The order is set to 10 which is arbitrary but should be lower priority than the default order of 0
         physx_interface = omni.physx.acquire_physx_interface()
         self._initialize_handle = physx_interface.get_simulation_event_stream_v2().create_subscription_to_pop_by_type(
             int(omni.physx.bindings._physx.SimulationEvent.RESUMED),
             lambda event, obj=weakref.proxy(self): obj._initialize_callback(event),
+            order=10,
         )
         self._invalidate_initialize_handle = (
             physx_interface.get_simulation_event_stream_v2().create_subscription_to_pop_by_type(
                 int(omni.physx.bindings._physx.SimulationEvent.STOPPED),
                 lambda event, obj=weakref.proxy(self): obj._invalidate_initialize_callback(event),
-            )
+                order=10,
+            ),
         )
         # add callback for debug visualization
         if self.cfg.debug_vis:

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

@@ -75,9 +75,8 @@ def height_scan(env: BaseEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
     """Height scan from the given sensor w.r.t. the sensor's frame."""
     # extract the used quantities (to enable type-hinting)
     sensor: RayCaster = env.scene.sensors[sensor_cfg.name]
-    # return the height scan
-    heights = sensor.data.pos_w[:, 2].unsqueeze(1) - sensor.data.ray_hits_w[..., 2] - 0.5
-    return heights
+    # height scan: height = sensor_height - hit_point_z - 0.5
+    return sensor.data.pos_w[:, 2].unsqueeze(1) - sensor.data.ray_hits_w[..., 2] - 0.5
 
 
 """

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/rl_env.py

@@ -329,10 +329,10 @@ class RLEnv(BaseEnv, gym.Env):
         """
         # update the curriculum for environments that need a reset
         self.curriculum_manager.compute(env_ids=env_ids)
-        # reset the internal buffers of the scene elements
-        self.scene.reset(env_ids)
         # randomize the MDP for environments that need a reset
         self.randomization_manager.randomize(env_ids=env_ids, mode="reset")
+        # reset the internal buffers of the scene elements
+        self.scene.reset(env_ids)
 
         # iterate over all managers and reset them
         # this returns a dictionary of information which is stored in the extras

source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/action_manager.py

@@ -215,7 +215,7 @@ class ActionManager(ManagerBase):
         """
         # check if action dimension is valid
         if self.total_action_dim != action.shape[1]:
-            raise ValueError(f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}")
+            raise ValueError(f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}.")
         # store the input actions
         self._prev_action[:] = self._action
         self._action[:] = action.to(self.device)

source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/observation_manager.py

@@ -123,7 +123,10 @@ class ObservationManager(ManagerBase):
         If a corruption/noise model is registered for a term, the function is called to corrupt
         the observation. The corruption function is expected to return a tensor with the same
         shape as the observation. The observations are clipped and scaled as per the configuration
-        settings. By default, no scaling or clipping is applied.
+        settings.
+
+        The operations are performed in the order: compute, add corruption/noise, clip, scale.
+        By default, no scaling or clipping is applied.
 
         Args:
             group_name: The name of the group for which to compute the observations. Defaults to :obj:`None`,

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/contact_sensor/contact_sensor.py

@@ -286,6 +286,10 @@ class ContactSensor(SensorBase):
         if self.contact_visualizer is None:
             visualizer_cfg = CONTACT_SENSOR_MARKER_CFG.replace(prim_path="/Visuals/ContactSensor")
             self.contact_visualizer = VisualizationMarkers(visualizer_cfg)
+        # safely return if view becomes invalid
+        # note: this invalidity happens because of isaac sim view callbacks
+        if self.body_physx_view is None:
+            return
         # marker indices
         # 0: contact, 1: no contact
         net_contact_force_w = torch.norm(self._data.net_forces_w, dim=-1)

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/sensor_base.py

@@ -53,16 +53,19 @@ class SensorBase(ABC):
 
         # note: Use weakref on callbacks to ensure that this object can be deleted when its destructor is called.
         # add callbacks for stage play/stop
+        # The order is set to 10 which is arbitrary but should be lower priority than the default order of 0
         physx_interface = omni.physx.acquire_physx_interface()
         self._initialize_handle = physx_interface.get_simulation_event_stream_v2().create_subscription_to_pop_by_type(
             int(omni.physx.bindings._physx.SimulationEvent.RESUMED),
             lambda event, obj=weakref.proxy(self): obj._initialize_callback(event),
+            order=10,
         )
         self._invalidate_initialize_handle = (
             physx_interface.get_simulation_event_stream_v2().create_subscription_to_pop_by_type(
                 int(omni.physx.bindings._physx.SimulationEvent.STOPPED),
                 lambda event, obj=weakref.proxy(self): obj._invalidate_initialize_callback(event),
-            )
+                order=10,
+            ),
         )
         # add callback for debug visualization
         if self.cfg.debug_vis:

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/simulation_context.py

@@ -66,7 +66,7 @@ class SimulationContext(_SimulationContext):
     .. _omni.isaac.core.simulation_context.SimulationContext: https://docs.omniverse.nvidia.com/py/isaacsim/source/extensions/omni.isaac.core/docs/index.html#module-omni.isaac.core.simulation_context
     """
 
-    class RenderMode(enum.Enum):
+    class RenderMode(enum.IntEnum):
         """Different rendering modes for the simulation.
 
         Render modes correspond to how the viewport and other UI elements (such as listeners to keyboard or mouse
@@ -294,8 +294,13 @@ class SimulationContext(_SimulationContext):
                 self._viewport_window.visible = False  # pyright: ignore [reportOptionalMemberAccess]
                 # reset the throttle counter
                 self._render_throttle_counter = 0
+            elif mode == self.RenderMode.NO_RENDERING:
+                # hide the viewport and disable updates
+                if self._viewport_context is not None:
+                    self._viewport_context.updates_enabled = False  # pyright: ignore [reportOptionalMemberAccess]
+                    self._viewport_window.visible = False  # pyright: ignore [reportOptionalMemberAccess]
             else:
-                raise ValueError(f"Unsupported rendering mode: {mode}. Supported modes are: {self.RenderMode}.")
+                raise ValueError(f"Unsupported render mode: {mode}! Please check `RenderMode` for details.")
             # update render mode
             self.render_mode = mode
 

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

@@ -13,7 +13,7 @@ from typing import Union
 
 import warp as wp
 
-__all__ = ["TENSOR_TYPES", "TENSOR_TYPE_CONVERSIONS", "convert_to_torch"]
+__all__ = ["TensorData", "TENSOR_TYPES", "TENSOR_TYPE_CONVERSIONS", "convert_to_torch"]
 
 TensorData = Union[np.ndarray, torch.Tensor, wp.array]
 """Type definition for a tensor data.

source/extensions/omni.isaac.orbit/omni/isaac/orbit/utils/warp/__init__.py

@@ -6,76 +6,9 @@
 
 """Operations based on warp."""
 
-import numpy as np
-import torch
+from .ops import convert_to_warp_mesh, raycast_mesh
 
-import warp as wp
-
-from .kernels import *
-
-__all__ = ["raycast_mesh", "convert_to_warp_mesh"]
-
-
-def raycast_mesh(ray_starts: torch.Tensor, ray_directions: torch.Tensor, mesh: wp.Mesh) -> torch.Tensor:
-    """Performs ray-casting against a mesh.
-
-    Note that the `ray_starts` and `ray_directions`, and `ray_hits` should have compatible shapes
-    and data types to ensure proper execution. Additionally, they all must be in the same frame.
-
-    Args:
-        ray_start: The starting position of the rays. Shape (N, 3).
-        ray_directions: The ray directions for each ray. Shape (N, 3).
-        mesh: The warp mesh to ray-cast against.
-
-    Returns:
-        The ray hit position. Shape (N, 3).
-            The returned tensor contains :obj:`float('inf')` for missed hits.
-    """
-    # extract device and shape information
-    shape = ray_starts.shape
-    device = ray_starts.device
-    # device of the mesh
-    mesh_device = wp.device_to_torch(mesh.device)
-    # reshape the tensors
-    ray_starts = ray_starts.to(mesh_device).view(-1, 3)
-    ray_directions = ray_directions.to(mesh_device).view(-1, 3)
-    num_rays = ray_starts.shape[0]
-    # create output tensor for the ray hits
-    ray_hits = torch.full((num_rays, 3), float("inf"), device=mesh_device)
-
-    # map the memory to warp arrays
-    ray_starts_wp = wp.from_torch(ray_starts, dtype=wp.vec3)
-    ray_directions_wp = wp.from_torch(ray_directions, dtype=wp.vec3)
-    ray_hits_wp = wp.from_torch(ray_hits, dtype=wp.vec3)
-
-    # launch the warp kernel
-    wp.launch(
-        kernel=raycast_mesh_kernel,
-        dim=num_rays,
-        inputs=[mesh.id, ray_starts_wp, ray_directions_wp, ray_hits_wp],
-        device=mesh.device,
-    )
-    # NOTE: Synchronize is not needed anymore, but we keep it for now. Check with @dhoeller.
-    wp.synchronize()
-
-    return ray_hits.to(device).view(shape)
-
-
-def convert_to_warp_mesh(points: np.ndarray, indices: np.ndarray, device: str) -> wp.Mesh:
-    """Create a warp mesh object with a mesh defined from vertices and triangles.
-
-    Args:
-        points: The vertices of the mesh. Shape is :math:`(N, 3)`, where :math:`N`
-            is the number of vertices.
-        indices: The triangles of the mesh as references to vertices for each triangle.
-            Shape is :math:`(M, 3)`, where :math:`M` is the number of triangles / faces.
-        device: The device to use for the mesh.
-
-    Returns:
-        The warp mesh object.
-    """
-    # create warp mesh
-    return wp.Mesh(
-        points=wp.array(points.astype(np.float32), dtype=wp.vec3, device=device),
-        indices=wp.array(indices.astype(np.int32).flatten(), dtype=wp.int32, device=device),
-    )
+__all__ = [
+    "raycast_mesh",
+    "convert_to_warp_mesh",
+]

source/extensions/omni.isaac.orbit/omni/isaac/orbit/utils/warp/ops.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+
+"""Wrapping around warp kernels for compatibility with torch tensors."""
+
+import numpy as np
+import torch
+
+import warp as wp
+
+from . import kernels
+
+
+def raycast_mesh(ray_starts: torch.Tensor, ray_directions: torch.Tensor, mesh: wp.Mesh) -> torch.Tensor:
+    """Performs ray-casting against a mesh.
+
+    Note that the `ray_starts` and `ray_directions`, and `ray_hits` should have compatible shapes
+    and data types to ensure proper execution. Additionally, they all must be in the same frame.
+
+    Args:
+        ray_start: The starting position of the rays. Shape (N, 3).
+        ray_directions: The ray directions for each ray. Shape (N, 3).
+        mesh: The warp mesh to ray-cast against.
+
+    Returns:
+        The ray hit position. Shape (N, 3). The returned tensor contains :obj:`float('inf')` for missed hits.
+    """
+    # extract device and shape information
+    shape = ray_starts.shape
+    device = ray_starts.device
+    # device of the mesh
+    mesh_device = wp.device_to_torch(mesh.device)
+    # reshape the tensors
+    ray_starts = ray_starts.to(mesh_device).view(-1, 3)
+    ray_directions = ray_directions.to(mesh_device).view(-1, 3)
+    num_rays = ray_starts.shape[0]
+    # create output tensor for the ray hits
+    ray_hits = torch.full((num_rays, 3), float("inf"), device=mesh_device)
+
+    # map the memory to warp arrays
+    ray_starts_wp = wp.from_torch(ray_starts, dtype=wp.vec3)
+    ray_directions_wp = wp.from_torch(ray_directions, dtype=wp.vec3)
+    ray_hits_wp = wp.from_torch(ray_hits, dtype=wp.vec3)
+
+    # launch the warp kernel
+    wp.launch(
+        kernel=kernels.raycast_mesh_kernel,
+        dim=num_rays,
+        inputs=[mesh.id, ray_starts_wp, ray_directions_wp, ray_hits_wp],
+        device=mesh.device,
+    )
+    # NOTE: Synchronize is not needed anymore, but we keep it for now. Check with @dhoeller.
+    wp.synchronize()
+
+    return ray_hits.to(device).view(shape)
+
+
+def convert_to_warp_mesh(points: np.ndarray, indices: np.ndarray, device: str) -> wp.Mesh:
+    """Create a warp mesh object with a mesh defined from vertices and triangles.
+
+    Args:
+        points: The vertices of the mesh. Shape is :math:`(N, 3)`, where :math:`N`
+            is the number of vertices.
+        indices: The triangles of the mesh as references to vertices for each triangle.
+            Shape is :math:`(M, 3)`, where :math:`M` is the number of triangles / faces.
+        device: The device to use for the mesh.
+
+    Returns:
+        The warp mesh object.
+    """
+    # create warp mesh
+    return wp.Mesh(
+        points=wp.array(points.astype(np.float32), dtype=wp.vec3, device=device),
+        indices=wp.array(indices.astype(np.int32).flatten(), dtype=wp.int32, device=device),
+    )

source/extensions/omni.isaac.orbit/setup.py

@@ -24,6 +24,10 @@ INSTALL_REQUIRES = [
     "tensordict",
     # devices
     "hidapi",
+    # gym
+    "gym==0.21.0",
+    "importlib-metadata~=4.13.0",
+    "setuptools<=66",  # setuptools 67.0 breaks gym
     # procedural-generation
     "trimesh",
     "pyglet==1.5.27",  # pyglet 2.0 requires python 3.8

source/extensions/omni.isaac.orbit/test/isaacsim/check_rep_texture_randomizer.py

@@ -9,6 +9,13 @@ Note:
     Currently this script fails since cloner does not support changing textures of cloned
     USD prims. This is because the prims are cloned using `Sdf.ChangeBlock` which does not
     allow individual texture changes.
+
+Usage:
+
+.. code-block:: bash
+
+    ./orbit.sh -p source/extensions/omni.isaac.orbit/test/isaacsim/check_rep_texture_randomizer.py
+
 """
 
 from __future__ import annotations
@@ -65,7 +72,7 @@ def main():
     set_camera_view([0.0, 30.0, 25.0], [0.0, 0.0, -2.5])
 
     # Parameters
-    num_balls = 2048
+    num_balls = 128
 
     # Create interface to clone the scene
     cloner = GridCloner(spacing=2.0)
@@ -81,7 +88,7 @@ def main():
     cloner.define_base_env("/World/envs")
     envs_prim_paths = cloner.generate_paths("/World/envs/env", num_paths=num_balls)
     env_positions = cloner.clone(
-        source_prim_path="/World/envs/env_0", prim_paths=envs_prim_paths, replicate_physics=True
+        source_prim_path="/World/envs/env_0", prim_paths=envs_prim_paths, replicate_physics=True, copy_from_source=True
     )
     physics_scene_path = sim.get_physics_context().prim_path
     cloner.filter_collisions(

source/extensions/omni.isaac.orbit/test/sensors/check_contact_sensor.py

@@ -100,7 +100,7 @@ def main():
     robot = Articulation(cfg=robot_cfg)
     # Contact sensor
     contact_sensor_cfg = ContactSensorCfg(
-        prim_path="/World/envs/env_.*/Robot/.*_SHANK", debug_vis=not args_cli.headless
+        prim_path="/World/envs/env_.*/Robot/.*_SHANK", track_air_time=True, debug_vis=not args_cli.headless
     )
     contact_sensor = ContactSensor(cfg=contact_sensor_cfg)
     # filter collisions within each environment instance

source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/locomotion/locomotion_env_cfg.py

@@ -117,6 +117,7 @@ class ObservationsCfg:
             func=mdp.height_scan,
             params={"sensor_cfg": SceneEntityCfg("height_scanner")},
             noise=Unoise(n_min=-0.1, n_max=0.1),
+            clip=(-1.0, 1.0),
         )
 
         def __post_init__(self):
@@ -254,7 +255,7 @@ class LocomotionEnvRoughCfg(RLEnvCfg):
     actions: ActionsCfg = ActionsCfg()
     commands: UniformVelocityCommandGeneratorCfg = UniformVelocityCommandGeneratorCfg(
         asset_name="robot",
-        resampling_time_range=(20.0, 20.0),
+        resampling_time_range=(10.0, 10.0),
         rel_standing_envs=0.02,
         rel_heading_envs=1.0,
         heading_command=True,

source/extensions/omni.isaac.orbit_envs/setup.py

@@ -23,10 +23,6 @@ INSTALL_REQUIRES = [
     "torch",
     "torchvision>=0.14.1",  # ensure compatibility with torch 1.13.1
     "protobuf==3.20.2",
-    # gym
-    "gym==0.21.0",
-    "importlib-metadata~=4.13.0",
-    "setuptools<=66",  # setuptools 67.0 breaks gym
     # data collection
     "h5py",
 ]

source/standalone/environments/list_envs.py

@@ -15,6 +15,18 @@ with `Isaac` in their name.
 
 from __future__ import annotations
 
+"""Launch Isaac Sim Simulator first."""
+
+
+from omni.isaac.orbit.app import AppLauncher
+
+# launch omniverse app
+app_launcher = AppLauncher(headless=True)
+simulation_app = app_launcher.app
+
+
+"""Rest everything follows."""
+
 import gym
 from prettytable import PrettyTable
 
@@ -46,4 +58,9 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    try:
+        # run the main function
+        main()
+    finally:
+        # close the app
+        simulation_app.close()