Fix backend asset state invalidation

source/isaaclab/changelog.d/backend-asset-invalidation.rst

@@ -0,0 +1,8 @@
+Added
+^^^^^
+
+* Added a ``skip_forward`` argument to the abstract root, body, and joint state writers of
+  :class:`~isaaclab.assets.articulation.base_articulation.BaseArticulation`,
+  :class:`~isaaclab.assets.rigid_object.base_rigid_object.BaseRigidObject`, and
+  :class:`~isaaclab.assets.rigid_object_collection.base_rigid_object_collection.BaseRigidObjectCollection`
+  to defer cached-buffer invalidation when several writes are batched before a single forward pass.

source/isaaclab/isaaclab/assets/articulation/base_articulation.py

@@ -345,6 +345,7 @@ def write_root_pose_to_sim_index(
         *,
         root_pose: torch.Tensor | wp.array,
         env_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root pose over selected environment indices into the simulation.
 
@@ -361,6 +362,7 @@ def write_root_pose_to_sim_index(
             root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7)
                 or (len(env_ids),) with dtype wp.transformf.
             env_ids: Environment indices. If None, then all indices are used.
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -370,6 +372,7 @@ def write_root_pose_to_sim_mask(
         *,
         root_pose: torch.Tensor | wp.array,
         env_mask: wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root pose over selected environment mask into the simulation.
 
@@ -386,6 +389,7 @@ def write_root_pose_to_sim_mask(
             root_pose: Root poses in simulation frame. Shape is (num_instances, 7)
                 or (num_instances,) with dtype wp.transformf.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -395,6 +399,7 @@ def write_root_link_pose_to_sim_index(
         *,
         root_pose: torch.Tensor | wp.array,
         env_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root link pose over selected environment indices into the simulation.
 
@@ -411,6 +416,7 @@ def write_root_link_pose_to_sim_index(
             root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7)
                 or (len(env_ids),) with dtype wp.transformf.
             env_ids: Environment indices. If None, then all indices are used.
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -420,6 +426,7 @@ def write_root_link_pose_to_sim_mask(
         *,
         root_pose: torch.Tensor | wp.array,
         env_mask: wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root link pose over selected environment mask into the simulation.
 
@@ -436,6 +443,7 @@ def write_root_link_pose_to_sim_mask(
             root_pose: Root poses in simulation frame. Shape is (num_instances, 7)
                 or (num_instances,) with dtype wp.transformf.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -445,6 +453,7 @@ def write_root_com_pose_to_sim_index(
         *,
         root_pose: torch.Tensor | wp.array,
         env_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root center of mass pose over selected environment indices into the simulation.
 
@@ -462,6 +471,7 @@ def write_root_com_pose_to_sim_index(
             root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7)
                 or (len(env_ids),) with dtype wp.transformf.
             env_ids: Environment indices. If None, then all indices are used.
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -471,6 +481,7 @@ def write_root_com_pose_to_sim_mask(
         *,
         root_pose: torch.Tensor | wp.array,
         env_mask: wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root center of mass pose over selected environment mask into the simulation.
 
@@ -488,6 +499,7 @@ def write_root_com_pose_to_sim_mask(
             root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7)
                 or (num_instances,) with dtype wp.transformf.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -497,6 +509,7 @@ def write_root_velocity_to_sim_index(
         *,
         root_velocity: torch.Tensor | wp.array,
         env_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root center of mass velocity over selected environment indices into the simulation.
 
@@ -516,6 +529,7 @@ def write_root_velocity_to_sim_index(
             root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6)
                 or (len(env_ids),) with dtype wp.spatial_vectorf.
             env_ids: Environment indices. If None, then all indices are used.
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -525,6 +539,7 @@ def write_root_velocity_to_sim_mask(
         *,
         root_velocity: torch.Tensor | wp.array,
         env_mask: wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root center of mass velocity over selected environment mask into the simulation.
 
@@ -544,6 +559,7 @@ def write_root_velocity_to_sim_mask(
             root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6)
                 or (num_instances,) with dtype wp.spatial_vectorf.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -553,6 +569,7 @@ def write_root_com_velocity_to_sim_index(
         *,
         root_velocity: torch.Tensor | wp.array,
         env_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root center of mass velocity over selected environment indices into the simulation.
 
@@ -572,6 +589,7 @@ def write_root_com_velocity_to_sim_index(
             root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6)
                 or (len(env_ids),) with dtype wp.spatial_vectorf.
             env_ids: Environment indices. If None, then all indices are used.
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -581,6 +599,7 @@ def write_root_com_velocity_to_sim_mask(
         *,
         root_velocity: torch.Tensor | wp.array,
         env_mask: wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root center of mass velocity over selected environment mask into the simulation.
 
@@ -600,6 +619,7 @@ def write_root_com_velocity_to_sim_mask(
             root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6)
                 or (num_instances,) with dtype wp.spatial_vectorf.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -609,6 +629,7 @@ def write_root_link_velocity_to_sim_index(
         *,
         root_velocity: torch.Tensor | wp.array,
         env_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root link velocity over selected environment indices into the simulation.
 
@@ -628,6 +649,7 @@ def write_root_link_velocity_to_sim_index(
             root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6)
                 or (len(env_ids),) with dtype wp.spatial_vectorf.
             env_ids: Environment indices. If None, then all indices are used.
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -637,6 +659,7 @@ def write_root_link_velocity_to_sim_mask(
         *,
         root_velocity: torch.Tensor | wp.array,
         env_mask: wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Set the root link velocity over selected environment mask into the simulation.
 
@@ -656,6 +679,7 @@ def write_root_link_velocity_to_sim_mask(
             root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6)
                 or (num_instances,) with dtype wp.spatial_vectorf.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -666,6 +690,7 @@ def write_joint_position_to_sim_index(
         position: torch.Tensor | wp.array,
         joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
         env_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Write joint positions to the simulation.
 
@@ -680,6 +705,7 @@ def write_joint_position_to_sim_index(
             position: Joint positions. Shape is (len(env_ids), len(joint_ids)).
             joint_ids: The joint indices to set the targets for. Defaults to None (all joints).
             env_ids: The environment indices to set the targets for. Defaults to None (all instances).
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -690,6 +716,7 @@ def write_joint_position_to_sim_mask(
         position: torch.Tensor | wp.array,
         joint_mask: wp.array | None = None,
         env_mask: wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Write joint positions to the simulation.
 
@@ -704,6 +731,7 @@ def write_joint_position_to_sim_mask(
             position: Joint positions. Shape is (num_instances, num_joints).
             joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,).
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -714,6 +742,7 @@ def write_joint_velocity_to_sim_index(
         velocity: torch.Tensor | wp.array,
         joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
         env_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Write joint velocities to the simulation.
 
@@ -728,6 +757,7 @@ def write_joint_velocity_to_sim_index(
             velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)).
             joint_ids: The joint indices to set the targets for. Defaults to None (all joints).
             env_ids: The environment indices to set the targets for. Defaults to None (all instances).
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()
 
@@ -738,6 +768,7 @@ def write_joint_velocity_to_sim_mask(
         velocity: torch.Tensor | wp.array,
         joint_mask: wp.array | None = None,
         env_mask: wp.array | None = None,
+        skip_forward: bool = False,
     ) -> None:
         """Write joint velocities to the simulation.
 
@@ -752,6 +783,7 @@ def write_joint_velocity_to_sim_mask(
             velocity: Joint velocities. Shape is (num_instances, num_joints).
             joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,).
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
+            skip_forward: Whether to skip invalidating cached data buffers after the write. Defaults to False.
         """
         raise NotImplementedError()