Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/solvers/rigid/rigid_solver_decomp.py (Python)

Code before:
                l_info = self.links_info[I_l]
                if l_info.n_dofs == 0:
                    continue

                dof_start = l_info.dof_start
                q_start = l_info.q_start
                q_end = l_info.q_end

                i_j = l_info.joint_start
                I_j = [i_j, i_b] if ti.static(self._options.batch_joints_info) else i_j
                joint_type = self.joints_info[I_j].type

                if joint_type == gs.JOINT_TYPE.FREE:
                    pos = ti.Vector([self.qpos[q_start, i_b], self.qpos[q_start + 1, i_b], self.qpos[q_start + 2, i_b]])
                    vel = ti.Vector(
                        [
                            self.dofs_state[dof_start, i_b].vel,
                            self.dofs_state[dof_start + 1, i_b].vel,
                            self.dofs_state[dof_start + 2, i_b].vel,
                        ]
                    )
                    pos = pos + vel * self._substep_dt
                    for j in ti.static(range(3)):
                        self.qpos[q_start + j, i_b] = pos[j]
                if joint_type == gs.JOINT_TYPE.SPHERICAL or joint_type == gs.JOINT_TYPE.FREE:
                    rot_offset = 3
                    if joint_type == gs.JOINT_TYPE.SPHERICAL:
                        rot_offset = 0
                    rot = ti.Vector(
                        [
                            self.qpos[q_start + rot_offset + 0, i_b],
                            self.qpos[q_start + rot_offset + 1, i_b],
                            self.qpos[q_start + rot_offset + 2, i_b],
                            self.qpos[q_start + rot_offset + 3, i_b],
                        ]
                    )
                    ang = (
                        ti.Vector(
                            [
                                self.dofs_state[dof_start + rot_offset + 0, i_b].vel,
                                self.dofs_state[dof_start + rot_offset + 1, i_b].vel,
                                self.dofs_state[dof_start + rot_offset + 2, i_b].vel,
                            ]
                        )
                        * self._substep_dt
                    )
                    qrot = gu.ti_rotvec_to_quat(ang)
                    rot = gu.ti_transform_quat_by_quat(qrot, rot)
                    for j in range(4):
                        self.qpos[q_start + j + rot_offset, i_b] = rot[j]
                else:

Review comment (suggestion):
Could one inlined ? Would look cleaner I would say: `rot_offset  = 3 if joint_type == gs.JOINT_TYPE.SPHERICAL else 0`

Code after:
                l_info = self.links_info[I_l]
                if l_info.n_dofs == 0:
                    continue

                dof_start = l_info.dof_start
                q_start = l_info.q_start
                q_end = l_info.q_end

                i_j = l_info.joint_start
                I_j = [i_j, i_b] if ti.static(self._options.batch_joints_info) else i_j
                joint_type = self.joints_info[I_j].type

                if joint_type == gs.JOINT_TYPE.FREE:
                    pos = ti.Vector([self.qpos[q_start, i_b], self.qpos[q_start + 1, i_b], self.qpos[q_start + 2, i_b]])
                    vel = ti.Vector(
                        [
                            self.dofs_state[dof_start, i_b].vel,
                            self.dofs_state[dof_start + 1, i_b].vel,
                            self.dofs_state[dof_start + 2, i_b].vel,
                        ]
                    )
                    pos = pos + vel * self._substep_dt
                    for j in ti.static(range(3)):
                        self.qpos[q_start + j, i_b] = pos[j]
                if joint_type == gs.JOINT_TYPE.SPHERICAL or joint_type == gs.JOINT_TYPE.FREE:
                    rot_offset = 3 if joint_type == gs.JOINT_TYPE.FREE else 0
                    rot = ti.Vector(
                        [
                            self.qpos[q_start + rot_offset + 0, i_b],
                            self.qpos[q_start + rot_offset + 1, i_b],
                            self.qpos[q_start + rot_offset + 2, i_b],
                            self.qpos[q_start + rot_offset + 3, i_b],
                        ]
                    )
                    ang = (
                        ti.Vector(
                            [
                                self.dofs_state[dof_start + rot_offset + 0, i_b].vel,
                                self.dofs_state[dof_start + rot_offset + 1, i_b].vel,
                                self.dofs_state[dof_start + rot_offset + 2, i_b].vel,
                            ]
                        )
                        * self._substep_dt
                    )
                    qrot = gu.ti_rotvec_to_quat(ang)
                    rot = gu.ti_transform_quat_by_quat(qrot, rot)
                    for j in range(4):
                        self.qpos[q_start + j + rot_offset, i_b] = rot[j]
                else:
                    for j in range(q_end - q_start):
                        self.qpos[q_start + j, i_b] = (