Fixes outdated sensor data after reset (#1276)

# Description

This change adds a call to `update_articulations_kinematic()` after
performing reset in an environment to ensure that non-render sensors are
updated after performing reset.

## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

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

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] 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
- [ ] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------
Signed-off-by: Kelly Guo <kellyguo123@hotmail.com>
Signed-off-by: Kelly Guo <kellyg@nvidia.com>

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

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.29.0"
+version = "0.29.1"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

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

 Changelog
 ---------
 
+0.29.1 (2024-12-15)
+~~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Added call to update articulation kinematics after reset to ensure states are updated for non-rendering sensors. Previously, some changes in reset such as modifying joint states would not be reflected in the rigid body states immediately after reset.
+
+
 0.29.0 (2024-12-15)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_rl_env.py

@@ -273,6 +273,10 @@ class DirectRLEnv(gym.Env):
         indices = torch.arange(self.num_envs, dtype=torch.int64, device=self.device)
         self._reset_idx(indices)
 
+        # update articulation kinematics
+        self.scene.write_data_to_sim()
+        self.sim.forward()
+
         # if sensors are added to the scene, make sure we render to reflect changes in reset
         if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
             self.sim.render()
@@ -346,6 +350,9 @@ class DirectRLEnv(gym.Env):
         reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1)
         if len(reset_env_ids) > 0:
             self._reset_idx(reset_env_ids)
+            # update articulation kinematics
+            self.scene.write_data_to_sim()
+            self.sim.forward()
             # if sensors are added to the scene, make sure we render to reflect changes in reset
             if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
                 self.sim.render()

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env.py

@@ -280,15 +280,17 @@ class ManagerBasedEnv:
 
         # reset state of scene
         self._reset_idx(env_ids)
-        self.scene.write_data_to_sim()
-
-        # trigger recorder terms for post-reset calls
-        self.recorder_manager.record_post_reset(env_ids)
 
+        # update articulation kinematics
+        self.scene.write_data_to_sim()
+        self.sim.forward()
         # if sensors are added to the scene, make sure we render to reflect changes in reset
         if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
             self.sim.render()
 
+        # trigger recorder terms for post-reset calls
+        self.recorder_manager.record_post_reset(env_ids)
+
         # compute observations
         self.obs_buf = self.observation_manager.compute()
 
@@ -328,13 +330,16 @@ class ManagerBasedEnv:
         # set the state
         self.scene.reset_to(state, env_ids, is_relative=is_relative)
 
-        # trigger recorder terms for post-reset calls
-        self.recorder_manager.record_post_reset(env_ids)
+        # update articulation kinematics
+        self.sim.forward()
 
         # if sensors are added to the scene, make sure we render to reflect changes in reset
         if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
             self.sim.render()
 
+        # trigger recorder terms for post-reset calls
+        self.recorder_manager.record_post_reset(env_ids)
+
         # compute observations
         self.obs_buf = self.observation_manager.compute()
 

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env.py

@@ -217,9 +217,9 @@ class ManagerBasedRLEnv(ManagerBasedEnv, gym.Env):
             self.recorder_manager.record_pre_reset(reset_env_ids)
 
             self._reset_idx(reset_env_ids)
-
-            # this is needed to make joint positions set from reset events effective
+            # update articulation kinematics
             self.scene.write_data_to_sim()
+            self.sim.forward()
 
             # if sensors are added to the scene, make sure we render to reflect changes in reset
             if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:

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

@@ -416,6 +416,14 @@ class SimulationContext(_SimulationContext):
         """
         return self._settings.get(name)
 
+    def forward(self) -> None:
+        """Updates articulation kinematics and fabric for rendering."""
+        if self._fabric_iface is not None:
+            if self.physics_sim_view is not None and self.is_playing():
+                # Update the articulations' link's poses before rendering
+                self.physics_sim_view.update_articulations_kinematic()
+            self._update_fabric(0.0, 0.0)
+
     """
     Operations - Override (standalone)
     """
@@ -486,11 +494,7 @@ class SimulationContext(_SimulationContext):
                 self.set_setting("/app/player/playSimulations", True)
         else:
             # manually flush the fabric data to update Hydra textures
-            if self._fabric_iface is not None:
-                if self.physics_sim_view is not None and self.is_playing():
-                    # Update the articulations' link's poses before rendering
-                    self.physics_sim_view.update_articulations_kinematic()
-                self._update_fabric(0.0, 0.0)
+            self.forward()
             # render the simulation
             # note: we don't call super().render() anymore because they do above operation inside
             #  and we don't want to do it twice. We may remove it once we drop support for Isaac Sim 2022.2.

source/extensions/omni.isaac.lab/test/sensors/test_outdated_sensor.py

+# Copyright (c) 2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Launch Isaac Sim Simulator first."""
+
+from omni.isaac.lab.app import AppLauncher, run_tests
+
+# launch the simulator
+app_launcher = AppLauncher(headless=True, enable_cameras=True)
+simulation_app = app_launcher.app
+
+
+"""Rest everything follows."""
+
+import gymnasium as gym
+import shutil
+import tempfile
+import torch
+import unittest
+
+import carb
+import omni.usd
+
+import omni.isaac.lab_tasks  # noqa: F401
+from omni.isaac.lab_tasks.utils.parse_cfg import parse_env_cfg
+
+
+class TestFrameTransformerAfterReset(unittest.TestCase):
+    """Test cases for checking FrameTransformer values after reset."""
+
+    @classmethod
+    def setUpClass(cls):
+        # this flag is necessary to prevent a bug where the simulation gets stuck randomly when running the
+        # test on many environments.
+        carb_settings_iface = carb.settings.get_settings()
+        carb_settings_iface.set_bool("/physics/cooking/ujitsoCollisionCooking", False)
+
+    def setUp(self):
+        # create a temporary directory to store the test datasets
+        self.temp_dir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        # delete the temporary directory after the test
+        shutil.rmtree(self.temp_dir)
+
+    def test_action_state_reocrder_terms(self):
+        """Check FrameTransformer values after reset."""
+        for task_name in ["Isaac-Stack-Cube-Franka-IK-Rel-v0"]:
+            for device in ["cuda:0", "cpu"]:
+                for num_envs in [1, 2]:
+                    with self.subTest(task_name=task_name, device=device):
+                        omni.usd.get_context().new_stage()
+
+                        # parse configuration
+                        env_cfg = parse_env_cfg(task_name, device=device, num_envs=num_envs)
+
+                        # create environment
+                        env = gym.make(task_name, cfg=env_cfg)
+
+                        # disable control on stop
+                        env.unwrapped.sim._app_control_on_stop_handle = None  # type: ignore
+
+                        # reset environment
+                        obs = env.reset()[0]
+
+                        # get the end effector position after the reset
+                        pre_reset_eef_pos = obs["policy"]["eef_pos"].clone()
+                        print(pre_reset_eef_pos)
+
+                        # step the environment with idle actions
+                        idle_actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device)
+                        obs = env.step(idle_actions)[0]
+
+                        # get the end effector position after the first step
+                        post_reset_eef_pos = obs["policy"]["eef_pos"]
+                        print(post_reset_eef_pos)
+
+                        # check if the end effector position is the same after the reset and the first step
+                        print(torch.all(torch.isclose(pre_reset_eef_pos, post_reset_eef_pos)))
+                        self.assertTrue(torch.all(torch.isclose(pre_reset_eef_pos, post_reset_eef_pos)))
+
+                        # close the environment
+                        env.close()
+
+
+if __name__ == "__main__":
+    run_tests()