Repository: Genesis-Embodied-AI/Genesis
File: genesis/ext/pyrender/interaction/mouse_spring.py (Python)

Code before:
        self.held_link = picked_link
        pose: Pose = Pose.from_link(self.held_link)
        self.held_point_in_local = pose.inverse_transform_point(control_point)
        self.prev_control_point = control_point

    def detach(self):
        self.held_link = None

    def apply_force(self, control_point: Vec3, delta_time: float):
        _ensure_torch_imported()

        # note when threaded: apply_force is called before attach!
        # note2: that was before we added a lock to ViewerInteraction; this migth be fixed now
        if not self.held_link:
            return
        
        self.prev_control_point = control_point

        # do simple force on COM only:
        link: RigidLink = self.held_link
        assert link
        lin_vel: Vec3 = Vec3.from_tensor(link.get_vel())
        ang_vel: Vec3 = Vec3.from_tensor(link.get_ang())
        link_pose: Pose = Pose.from_link(link)
        held_point_in_world: Vec3 = link_pose.transform_point(self.held_point_in_local)

        # note: we should assert earlier that link inertial_pos/quat are not None
        # todo: verify inertial_pos/quat are stored in local frame
        link_T_principal: Pose = Pose(Vec3.from_arraylike(link.inertial_pos), Quat.from_arraylike(link.inertial_quat))
        world_T_principal: Pose = link_pose * link_T_principal

        arm_in_principal: Vec3 = link_T_principal.get_inverse() * self.held_point_in_local   # for non-spherical inertia
        arm_in_world: Vec3 = world_T_principal.rot * arm_in_principal  # for spherical inertia

        pos_err_v: Vec3 = control_point - held_point_in_world
        inv_mass: float = float(1.0 / link.get_mass() if link.get_mass() > 0.0 else 0.0)
        inv_spherical_inertia: float = float(1.0 / link.inertial_i[0, 0] if link.inertial_i[0, 0] > 0.0 else 0.0)

        inv_dt: float = 1.0 / delta_time
        tau: float = MOUSE_SPRING_POSITION_CORRECTION_FACTOR
        damp: float = MOUSE_SPRING_VELOCITY_CORRECTION_FACTOR

        total_impulse: Vec3 = Vec3.zero()
        total_torque_impulse: Vec3 = Vec3.zero()

        for i in range(3*4):
            body_point_vel: Vec3 = lin_vel + ang_vel.cross(arm_in_world)
            vel_err_v: Vec3 = Vec3.zero() - body_point_vel

            dir: Vec3 = Vec3.zero()
            dir.v[i % 3] = 1.0

Review comment (suggestion):
If the simulation is built, `inertial_pos` and `inertial_quat` should not be none. We should update these property getters to read values from taichi field if the simulation is already built. See `sol_params` implementation: https://github.com/Genesis-Embodied-AI/Genesis/blob/767d536b8689cba5f382fe184df8bd270f889ea2/genesis/engine/entities/rigid_entity/rigid_joint.py#L145

Code after:
        self.held_link = picked_link
        pose: Pose = Pose.from_link(self.held_link)
        self.held_point_in_local = pose.inverse_transform_point(control_point)
        self.prev_control_point = control_point

    def detach(self) -> None:
        self.held_link = None

    def apply_force(self, control_point: Vec3, delta_time: float) -> None:
        _ensure_torch_imported()

        # note when threaded: apply_force is called before attach!
        # note2: that was before we added a lock to ViewerInteraction; this migth be fixed now
        if not self.held_link:
            return
        
        self.prev_control_point = control_point

        # do simple force on COM only:
        link: RigidLink = self.held_link
        lin_vel: Vec3 = Vec3.from_tensor(link.get_vel())
        ang_vel: Vec3 = Vec3.from_tensor(link.get_ang())
        link_pose: Pose = Pose.from_link(link)
        held_point_in_world: Vec3 = link_pose.transform_point(self.held_point_in_local)

        # note: we should assert earlier that link inertial_pos/quat are not None
        # todo: verify inertial_pos/quat are stored in local frame
        link_T_principal: Pose = Pose(Vec3.from_arraylike(link.inertial_pos), Quat.from_arraylike(link.inertial_quat))
        world_T_principal: Pose = link_pose * link_T_principal

        arm_in_principal: Vec3 = link_T_principal.inverse_transform_point(self.held_point_in_local)   # for non-spherical inertia
        arm_in_world: Vec3 = world_T_principal.rot * arm_in_principal  # for spherical inertia

        pos_err_v: Vec3 = control_point - held_point_in_world
        inv_mass: float = float(1.0 / link.get_mass() if link.get_mass() > 0.0 else 0.0)
        inv_spherical_inertia: float = float(1.0 / link.inertial_i[0, 0] if link.inertial_i[0, 0] > 0.0 else 0.0)

        inv_dt: float = 1.0 / delta_time
        tau: float = MOUSE_SPRING_POSITION_CORRECTION_FACTOR
        damp: float = MOUSE_SPRING_VELOCITY_CORRECTION_FACTOR

        total_impulse: Vec3 = Vec3.zero()
        total_torque_impulse: Vec3 = Vec3.zero()

        for i in range(3*4):
            body_point_vel: Vec3 = lin_vel + ang_vel.cross(arm_in_world)
            vel_err_v: Vec3 = Vec3.zero() - body_point_vel

            dir: Vec3 = Vec3.zero()
            dir.v[i % 3] = 1.0
            pos_err: float = dir.dot(pos_err_v)