From 5ae1cd390cf3638c4c5eb068f2619f22d779c465 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Mon, 16 Mar 2026 18:37:13 +0100 Subject: [PATCH] Add lazy velocity update to kinematic entity. --- genesis/engine/solvers/kinematic_solver.py | 39 ++++++++++++++++--- .../solvers/rigid/abd/forward_kinematics.py | 22 ----------- 2 files changed, 33 insertions(+), 28 deletions(-) diff --git a/genesis/engine/solvers/kinematic_solver.py b/genesis/engine/solvers/kinematic_solver.py index 93c1c2845..f775a9feb 100644 --- a/genesis/engine/solvers/kinematic_solver.py +++ b/genesis/engine/solvers/kinematic_solver.py @@ -446,8 +446,23 @@ def substep_pre_coupling(self, f): def substep_pre_coupling_grad(self, f): pass + def _ensure_forward_vel_updated(self): + if not self._is_forward_vel_updated: + kernel_forward_velocity( + self.scene._envs_idx, + links_state=self.links_state, + links_info=self.links_info, + joints_info=self.joints_info, + dofs_state=self.dofs_state, + entities_info=self.entities_info, + rigid_global_info=self._rigid_global_info, + static_rigid_sim_config=self._static_rigid_sim_config, + is_backward=False, + ) + self._is_forward_vel_updated = True + def substep_post_coupling(self, f): - if not self._is_forward_pos_updated or not self._is_forward_vel_updated: + if not self._is_forward_pos_updated: kernel_forward_kinematics( self.scene._envs_idx, links_state=self.links_state, @@ -461,7 +476,6 @@ def substep_post_coupling(self, f): static_rigid_sim_config=self._static_rigid_sim_config, ) self._is_forward_pos_updated = True - self._is_forward_vel_updated = True def substep_post_coupling_grad(self, f): pass @@ -578,6 +592,17 @@ def set_state(self, f, state, envs_idx=None): rigid_global_info=self._rigid_global_info, static_rigid_sim_config=self._static_rigid_sim_config, ) + kernel_forward_velocity( + envs_idx, + links_state=self.links_state, + links_info=self.links_info, + joints_info=self.joints_info, + dofs_state=self.dofs_state, + entities_info=self.entities_info, + rigid_global_info=self._rigid_global_info, + static_rigid_sim_config=self._static_rigid_sim_config, + is_backward=False, + ) self._is_forward_pos_updated = True self._is_forward_vel_updated = True @@ -680,7 +705,7 @@ def set_base_links_pos(self, pos, links_idx=None, envs_idx=None, *, relative=Fal static_rigid_sim_config=self._static_rigid_sim_config, ) self._is_forward_pos_updated = True - self._is_forward_vel_updated = True + self._is_forward_vel_updated = False def set_base_links_pos_grad(self, links_idx, envs_idx, relative, pos_grad): if links_idx is None: @@ -734,7 +759,7 @@ def set_base_links_quat(self, quat, links_idx=None, envs_idx=None, *, relative=F static_rigid_sim_config=self._static_rigid_sim_config, ) self._is_forward_pos_updated = True - self._is_forward_vel_updated = True + self._is_forward_vel_updated = False def set_base_links_quat_grad(self, links_idx, envs_idx, relative, quat_grad): if links_idx is None: @@ -805,7 +830,7 @@ def set_qpos(self, qpos, qs_idx=None, envs_idx=None, *, skip_forward=False): static_rigid_sim_config=self._static_rigid_sim_config, ) self._is_forward_pos_updated = True - self._is_forward_vel_updated = True + self._is_forward_vel_updated = False else: self._is_forward_pos_updated = False self._is_forward_vel_updated = False @@ -917,7 +942,7 @@ def set_dofs_position(self, position, dofs_idx=None, envs_idx=None): static_rigid_sim_config=self._static_rigid_sim_config, ) self._is_forward_pos_updated = True - self._is_forward_vel_updated = True + self._is_forward_vel_updated = False def get_links_pos(self, links_idx=None, envs_idx=None): if not gs.use_zerocopy: @@ -932,6 +957,7 @@ def get_links_quat(self, links_idx=None, envs_idx=None): return tensor[0] if self.n_envs == 0 else tensor def get_links_vel(self, links_idx=None, envs_idx=None): + self._ensure_forward_vel_updated() if gs.use_zerocopy: mask = (0, *indices_to_mask(links_idx)) if self.n_envs == 0 else indices_to_mask(envs_idx, links_idx) cd_vel = qd_to_torch(self.links_state.cd_vel, transpose=True) @@ -949,6 +975,7 @@ def get_links_vel(self, links_idx=None, envs_idx=None): return _tensor def get_links_ang(self, links_idx=None, envs_idx=None): + self._ensure_forward_vel_updated() tensor = qd_to_torch(self.links_state.cd_ang, envs_idx, links_idx, transpose=True, copy=True) return tensor[0] if self.n_envs == 0 else tensor diff --git a/genesis/engine/solvers/rigid/abd/forward_kinematics.py b/genesis/engine/solvers/rigid/abd/forward_kinematics.py index 3c5dc0339..7f1166981 100644 --- a/genesis/engine/solvers/rigid/abd/forward_kinematics.py +++ b/genesis/engine/solvers/rigid/abd/forward_kinematics.py @@ -157,17 +157,6 @@ def kernel_forward_kinematics( static_rigid_sim_config=static_rigid_sim_config, is_backward=False, ) - func_forward_velocity_batch( - i_b=i_b, - entities_info=entities_info, - links_info=links_info, - links_state=links_state, - joints_info=joints_info, - dofs_state=dofs_state, - rigid_global_info=rigid_global_info, - static_rigid_sim_config=static_rigid_sim_config, - is_backward=False, - ) @qd.kernel(fastcache=gs.use_fastcache) @@ -211,17 +200,6 @@ def kernel_masked_forward_kinematics( static_rigid_sim_config=static_rigid_sim_config, is_backward=False, ) - func_forward_velocity_batch( - i_b=i_b, - entities_info=entities_info, - links_info=links_info, - links_state=links_state, - joints_info=joints_info, - dofs_state=dofs_state, - rigid_global_info=rigid_global_info, - static_rigid_sim_config=static_rigid_sim_config, - is_backward=False, - ) @qd.kernel(fastcache=gs.use_fastcache)