def snap_rotate_face_with_threshold()

in robogym/envs/dactyl/common/cube_manipulator.py [0:0]

• 27 lines of code
• 8 McCabe index (conditional complexity)

    def snap_rotate_face_with_threshold(self, axis, side, angle, threshold=0.1):
        """
        Rotate face of a cube in a "snapping" fashion, correcting the cube along the way.
        Underlying assumption: cube is already in a "snapped", physically-aligned state

        Threshold is threshold in radians which decides maximum angle we want to snap over.
        If the angle required to move the face to be snapped is larger than that, the cube
        will remain locked and won't rotate.
        """
        qpos = self.sim.data.qpos

        drivers = rotation.normalize_angles(
            qpos[[self.joints_qpos_map[x] for x in self.drivers]]
        )

        perpendicular_axes = sorted({0, 1, 2} - {axis})

        transaction = []
        abort = False

        for other_axis in perpendicular_axes:
            for other_side in range(2):
                other_driver_idx = other_axis * 2 + other_side

                other_angle = drivers[other_driver_idx]
                other_angle_aligned = rotation.round_to_straight_angles(other_angle)
                other_angle_diff = rotation.normalize_angles(
                    other_angle_aligned - other_angle
                )

                if (
                    np.abs(other_angle_diff) < np.abs(angle)
                    and np.abs(other_angle_diff) < threshold
                ):
                    transaction.append((other_axis, other_side, other_angle_diff))
                else:
                    abort = True

        if not abort:
            # Snap other faces
            for other_axis, other_side, angle_diff in transaction:
                self.rotate_face(other_axis, other_side, angle_diff)

            # rotate the actual face
            self.rotate_face(axis, side, angle)