From c6af307a2b1a62432610595672fbadfdedc22a7b Mon Sep 17 00:00:00 2001
From: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com>
Date: Fri, 20 Oct 2023 15:11:00 -0400
Subject: [PATCH] Fix the double ellipsis issue when resolving indices (#196)

# Description

Earlier in the code, we were using the `Ellipsis` object to index the
dimensions of the tensor. This led to situations where
we indexed multi-dimension tensors as: `x[..., ..., 0]`. This now leads
to errors with Python 3.10.

The MR replaces `Ellipsis` with the `slice(None)` object, which results
in indexing as: `x[:, :, 0]`.

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

---------

Co-authored-by: Farbod Farshidian <ffarshidian@theaiinstitute.com>
---
 .../omni.isaac.orbit/config/extension.toml    |  2 +-
 .../omni.isaac.orbit/docs/CHANGELOG.rst       | 10 +++++
 .../isaac/orbit/actuators/actuator_base.py    | 13 ++++---
 .../orbit/assets/articulation/articulation.py | 36 +++++++++---------
 .../orbit/assets/rigid_object/rigid_object.py |  6 +--
 .../command_generator_base.py                 |  2 +-
 .../orbit/envs/mdp/actions/joint_actions.py   |  2 +-
 .../envs/mdp/actions/task_space_actions.py    |  2 +-
 .../isaac/orbit/managers/action_manager.py    |  2 +-
 .../orbit/managers/curriculum_manager.py      |  2 +-
 .../isaac/orbit/managers/reward_manager.py    |  2 +-
 .../orbit/managers/termination_manager.py     |  2 +-
 .../sensors/contact_sensor/contact_sensor.py  |  6 +--
 .../omni/isaac/orbit/sensors/sensor_base.py   |  2 +-
 .../isaac/orbit/sim/simulation_context.py     |  3 ++
 .../sim/spawners/from_files/from_files.py     |  2 +-
 .../test/isaacsim/test_torch.py               | 38 +++++++++++++++++++
 .../environments/state_machine/play_lift.py   |  2 +-
 18 files changed, 93 insertions(+), 41 deletions(-)
 create mode 100644 source/extensions/omni.isaac.orbit/test/isaacsim/test_torch.py

diff --git a/source/extensions/omni.isaac.orbit/config/extension.toml b/source/extensions/omni.isaac.orbit/config/extension.toml
index f6129b54cc..0509171218 100644
--- a/source/extensions/omni.isaac.orbit/config/extension.toml
+++ b/source/extensions/omni.isaac.orbit/config/extension.toml
@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.9.12"
+version = "0.9.13"
 
 # Description
 title = "ORBIT framework for Robot Learning"
diff --git a/source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst b/source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst
index cf5c1a2d89..389c246928 100644
--- a/source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst
+++ b/source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst
@@ -1,6 +1,16 @@
 Changelog
 ---------
 
+0.9.13 (2023-10-20)
+~~~~~~~~~~~~~~~~~~~
+
+Fixed
+^^^^^
+
+* Fixed the issue with double :obj:`Ellipsis` when indexing tensors with multiple dimensions.
+  The fix now uses :obj:`slice(None)` instead of :obj:`Ellipsis` to index the tensors.
+
+
 0.9.12 (2023-10-18)
 ~~~~~~~~~~~~~~~~~~~
 
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/actuators/actuator_base.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/actuators/actuator_base.py
index 0146c46101..b1cf27e7c5 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/actuators/actuator_base.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/actuators/actuator_base.py
@@ -45,15 +45,15 @@ class ActuatorBase(ABC):
     """The damping (D gain) of the PD controller. Shape is ``(num_envs, num_joints)``."""
 
     def __init__(
-        self, cfg: ActuatorBaseCfg, joint_names: list[str], joint_ids: list[int] | Ellipsis, num_envs: int, device: str
+        self, cfg: ActuatorBaseCfg, joint_names: list[str], joint_ids: Sequence[int], num_envs: int, device: str
     ):
         """Initialize the actuator.
 
         Args:
             cfg: The configuration of the actuator model.
             joint_names: The joint names in the articulation.
-            joint_ids: The joint indices in the articulation. If :obj:`Ellipsis`, then all the joints in the
-                articulation are part of the group.
+            joint_ids: The joint indices in the articulation. If :obj:`slice(None)`, then all
+                the joints in the articulation are part of the group.
             num_envs: Number of articulations in the view.
             device: Device used for processing.
         """
@@ -99,7 +99,7 @@ def __str__(self) -> str:
         """Returns: A string representation of the actuator group."""
         # resolve joint indices for printing
         joint_indices = self.joint_indices
-        if joint_indices is Ellipsis:
+        if joint_indices == slice(None):
             joint_indices = list(range(self.num_joints))
         return (
             f"<class {self.__class__.__name__}> object:\n"
@@ -124,11 +124,12 @@ def joint_names(self) -> list[str]:
         return self._joint_names
 
     @property
-    def joint_indices(self) -> list[int] | Ellipsis:
+    def joint_indices(self) -> Sequence[int]:
         """Articulation's joint indices that are part of the group.
 
         Note:
-            If :obj:`Ellipsis` is returned, then the group contains all the joints in the articulation.
+            If :obj:`slice(None)` is returned, then the group contains all the joints in the articulation.
+            We do this to avoid unnecessary indexing of the joints for performance reasons.
         """
         return self._joint_indices
 
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/articulation/articulation.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/articulation/articulation.py
index 8765ecc7cb..1e722c81ae 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/articulation/articulation.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/articulation/articulation.py
@@ -101,7 +101,7 @@ def reset(self, env_ids: Sequence[int] | None = None):
         super().reset(env_ids)
         # use ellipses object to skip initial indices.
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # reset actuators
         for actuator in self.actuators.values():
             actuator.reset(env_ids)
@@ -162,7 +162,7 @@ def find_joints(
     def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None):
         # resolve all indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # note: we need to do this here since tensors are not set into simulation until step.
         # set into internal buffers
         self._data.root_state_w[env_ids, :7] = root_pose.clone()
@@ -175,7 +175,7 @@ def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int]
     def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None):
         # resolve all indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # note: we need to do this here since tensors are not set into simulation until step.
         # set into internal buffers
         self._data.root_state_w[env_ids, 7:] = root_velocity.clone()
@@ -199,9 +199,9 @@ def write_joint_state_to_sim(
         """
         # resolve indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         if joint_ids is None:
-            joint_ids = ...
+            joint_ids = slice(None)
         # set into internal buffers
         self._data.joint_pos[env_ids, joint_ids] = position
         self._data.joint_vel[env_ids, joint_ids] = velocity
@@ -227,9 +227,9 @@ def write_joint_stiffness_to_sim(
         # note: This function isn't setting the values for actuator models. (#128)
         # resolve indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         if joint_ids is None:
-            joint_ids = ...
+            joint_ids = slice(None)
         # set into internal buffers
         self._data.joint_stiffness[env_ids, joint_ids] = stiffness
         # set into simulation
@@ -251,9 +251,9 @@ def write_joint_damping_to_sim(
         # note: This function isn't setting the values for actuator models. (#128)
         # resolve indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         if joint_ids is None:
-            joint_ids = ...
+            joint_ids = slice(None)
         # set into internal buffers
         self._data.joint_damping[env_ids, joint_ids] = damping
         # set into simulation
@@ -276,9 +276,9 @@ def write_joint_torque_limit_to_sim(
         # note: This function isn't setting the values for actuator models. (#128)
         # resolve indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         if joint_ids is None:
-            joint_ids = ...
+            joint_ids = slice(None)
         # set into internal buffers
         torque_limit_all = self.root_physx_view.get_dof_max_forces()
         torque_limit_all[env_ids, joint_ids] = limits
@@ -306,9 +306,9 @@ def set_joint_position_target(
         """
         # resolve indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         if joint_ids is None:
-            joint_ids = ...
+            joint_ids = slice(None)
         # set targets
         self._data.joint_pos_target[env_ids, joint_ids] = target
 
@@ -328,9 +328,9 @@ def set_joint_velocity_target(
         """
         # resolve indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         if joint_ids is None:
-            joint_ids = ...
+            joint_ids = slice(None)
         # set targets
         self._data.joint_vel_target[env_ids, joint_ids] = target
 
@@ -350,9 +350,9 @@ def set_joint_effort_target(
         """
         # resolve indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         if joint_ids is None:
-            joint_ids = ...
+            joint_ids = slice(None)
         # set targets
         self._data.joint_effort_target[env_ids, joint_ids] = target
 
@@ -513,7 +513,7 @@ def _process_actuators_cfg(self):
                 )
             # for efficiency avoid indexing when over all indices
             if len(joint_names) == self.num_joints:
-                joint_ids = ...
+                joint_ids = slice(None)
             # create actuator collection
             actuator: ActuatorBase = actuator_cfg.class_type(
                 cfg=actuator_cfg,
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/rigid_object/rigid_object.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/rigid_object/rigid_object.py
index f0153a2fd7..2aa63cf153 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/rigid_object/rigid_object.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/rigid_object/rigid_object.py
@@ -100,7 +100,7 @@ def body_physx_view(self) -> physx.RigidBodyView:
     def reset(self, env_ids: Sequence[int] | None = None):
         # resolve all indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # reset external wrench
         self._external_force_b[env_ids] = 0.0
         self._external_torque_b[env_ids] = 0.0
@@ -171,7 +171,7 @@ def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int]
         """
         # resolve all indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # note: we need to do this here since tensors are not set into simulation until step.
         # set into internal buffers
         self._data.root_state_w[env_ids, :7] = root_pose.clone()
@@ -190,7 +190,7 @@ def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Seque
         """
         # resolve all indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # note: we need to do this here since tensors are not set into simulation until step.
         # set into internal buffers
         self._data.root_state_w[env_ids, 7:] = root_velocity.clone()
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/command_generators/command_generator_base.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/command_generators/command_generator_base.py
index ba819ab3ba..e1de4b5a6f 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/command_generators/command_generator_base.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/command_generators/command_generator_base.py
@@ -122,7 +122,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]:
         """
         # resolve the environment IDs
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # set the command counter to zero
         self.command_counter[env_ids] = 0
         # resample the command
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/joint_actions.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/joint_actions.py
index 4c414b2263..a1770dc6fc 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/joint_actions.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/joint_actions.py
@@ -65,7 +65,7 @@ def __init__(self, cfg: actions_cfg.JointActionCfg, env: BaseEnv) -> None:
 
         # Avoid indexing across all joints for efficiency
         if self._num_joints == self._asset.num_joints:
-            self._joint_ids = ...
+            self._joint_ids = slice(None)
 
         # create tensors for raw and processed actions
         self._raw_actions = torch.zeros(self.num_envs, self.action_dim, device=self.device)
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/task_space_actions.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/task_space_actions.py
index 645a891eae..12fa634bb2 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/task_space_actions.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/task_space_actions.py
@@ -76,7 +76,7 @@ def __init__(self, cfg: actions_cfg.DifferentialInverseKinematicsActionCfg, env:
         )
         # Avoid indexing across all joints for efficiency
         if self._num_joints == self._asset.num_joints:
-            self._joint_ids = ...
+            self._joint_ids = slice(None)
 
         # create the differential IK controller
         self._controller = DifferentialIKController(cfg=self.cfg.controller, num_envs=self.num_envs, device=self.device)
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/action_manager.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/action_manager.py
index 4b26d86793..034bd11a75 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/action_manager.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/action_manager.py
@@ -197,7 +197,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]
         """
         # resolve environment ids
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # reset the action history
         self._prev_action[env_ids] = 0.0
         self._action[env_ids] = 0.0
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/curriculum_manager.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/curriculum_manager.py
index cbabc3e863..d5e00f2915 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/curriculum_manager.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/curriculum_manager.py
@@ -120,7 +120,7 @@ def compute(self, env_ids: Sequence[int] | None = None):
         """
         # resolve environment indices
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # iterate over all the curriculum terms
         for name, term_cfg in zip(self._term_names, self._term_cfgs):
             state = term_cfg.func(self._env, env_ids, **term_cfg.params)
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/reward_manager.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/reward_manager.py
index b2328580fa..fb1b34eddf 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/reward_manager.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/reward_manager.py
@@ -99,7 +99,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]
         """
         # resolve environment ids
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # store information
         extras = {}
         for key in self._episode_sums.keys():
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/termination_manager.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/termination_manager.py
index 361730b303..ef2e2af3f1 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/termination_manager.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/termination_manager.py
@@ -103,7 +103,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]
         """
         # resolve environment ids
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # add to episode dict
         extras = {}
         for key in self._episode_dones.keys():
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/contact_sensor/contact_sensor.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/contact_sensor/contact_sensor.py
index 7757135ea9..3aa4fc6764 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/contact_sensor/contact_sensor.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/contact_sensor/contact_sensor.py
@@ -135,7 +135,7 @@ def reset(self, env_ids: Sequence[int] | None = None):
         super().reset(env_ids)
         # resolve None
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # reset accumulative data buffers
         self._data.current_air_time[env_ids] = 0.0
         self._data.last_air_time[env_ids] = 0.0
@@ -223,7 +223,7 @@ def _update_buffers_impl(self, env_ids: Sequence[int]):
         """Fills the buffers of the sensor data."""
         # default to all sensors
         if len(env_ids) == self._num_envs:
-            env_ids = ...
+            env_ids = slice(None)
         # obtain the poses of the sensors:
         # TODO decide if we really to track poses -- This is the body's CoM. Not contact location.
         pose = self.body_physx_view.get_transforms()
@@ -257,7 +257,7 @@ def _update_buffers_impl(self, env_ids: Sequence[int]):
         # since this function is called every frame, we can use the difference to get the elapsed time
         elapsed_time = self._timestamp[env_ids] - self._timestamp_last_update[env_ids]
         # -- check contact state of bodies
-        is_contact = torch.norm(self._data.net_forces_w[env_ids, 0, :, :], dim=-1) > 1.0
+        is_contact = torch.norm(self._data.net_forces_w[env_ids, :, :], dim=-1) > 1.0
         is_first_contact = (self._data.current_air_time[env_ids] > 0) * is_contact
         # -- update ongoing timer for bodies air
         self._data.current_air_time[env_ids] += elapsed_time.unsqueeze(-1)
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/sensor_base.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/sensor_base.py
index f1256f5a4a..79733e3112 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/sensor_base.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/sensor_base.py
@@ -140,7 +140,7 @@ def reset(self, env_ids: Sequence[int] | None = None):
         """
         # Resolve sensor ids
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         # Reset the timestamp for the sensors
         self._timestamp[env_ids] = 0.0
         self._timestamp_last_update[env_ids] = 0.0
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/simulation_context.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/simulation_context.py
index 8608a5c15f..e4d2ad85b0 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/simulation_context.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/simulation_context.py
@@ -380,6 +380,9 @@ def render(self, mode: RenderMode | None = None):
                 self._render_throttle_counter = 0
                 # here we don't render viewport so don't need to flush flatcache
                 super().render()
+        elif self.render_mode == self.RenderMode.HEADLESS:
+            # we never want to render anything here -- camera rendering will also not work then?
+            pass
         else:
             # this is called even if we are in headless mode - we allow this for off-screen rendering
             # manually flush the flatcache data to update Hydra textures
diff --git a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/spawners/from_files/from_files.py b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/spawners/from_files/from_files.py
index 6d694e72b3..b3de7febeb 100644
--- a/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/spawners/from_files/from_files.py
+++ b/source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/spawners/from_files/from_files.py
@@ -9,8 +9,8 @@
 
 import carb
 import omni.isaac.core.utils.prims as prim_utils
-from omni.isaac.version import get_version
 import omni.kit.commands
+from omni.isaac.version import get_version
 from pxr import Gf, Sdf, Usd
 
 from omni.isaac.orbit.sim import loaders, schemas
diff --git a/source/extensions/omni.isaac.orbit/test/isaacsim/test_torch.py b/source/extensions/omni.isaac.orbit/test/isaacsim/test_torch.py
new file mode 100644
index 0000000000..94e6e01be7
--- /dev/null
+++ b/source/extensions/omni.isaac.orbit/test/isaacsim/test_torch.py
@@ -0,0 +1,38 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+import unittest
+
+
+class TestTorchOperations(unittest.TestCase):
+    """Tests for assuring torch related operations used in Orbit."""
+
+    def test_array_slicing(self):
+        """Check that using ellipsis and slices work for torch tensors."""
+
+        size = (400, 300, 5)
+        my_tensor = torch.rand(size, device="cuda:0")
+
+        self.assertEqual(my_tensor[..., 0].shape, (400, 300))
+        self.assertEqual(my_tensor[:, :, 0].shape, (400, 300))
+        self.assertEqual(my_tensor[slice(None), slice(None), 0].shape, (400, 300))
+        with self.assertRaises(IndexError):
+            my_tensor[..., ..., 0]
+
+        self.assertEqual(my_tensor[0, ...].shape, (300, 5))
+        self.assertEqual(my_tensor[0, :, :].shape, (300, 5))
+        self.assertEqual(my_tensor[0, slice(None), slice(None)].shape, (300, 5))
+        self.assertEqual(my_tensor[0, ..., ...].shape, (300, 5))
+
+        self.assertEqual(my_tensor[..., 0, 0].shape, (400,))
+        self.assertEqual(my_tensor[slice(None), 0, 0].shape, (400,))
+        self.assertEqual(my_tensor[:, 0, 0].shape, (400,))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/standalone/environments/state_machine/play_lift.py b/source/standalone/environments/state_machine/play_lift.py
index 852a1ea4ce..584283c065 100644
--- a/source/standalone/environments/state_machine/play_lift.py
+++ b/source/standalone/environments/state_machine/play_lift.py
@@ -201,7 +201,7 @@ def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu")
     def reset_idx(self, env_ids: Sequence[int] = None):
         """Reset the state machine."""
         if env_ids is None:
-            env_ids = ...
+            env_ids = slice(None)
         self.sm_state[env_ids] = 0
         self.sm_wait_time[env_ids] = 0.0