def rotate_face()

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

• 26 lines of code
• 7 McCabe index (conditional complexity)

    def rotate_face(self, axis, side, angle):
        """
        Rotate given face (identified by axis and side) by given angle.
        Cube should be in a reasonably aligned state for this to work well.
        """
        assert 0 <= axis <= 2
        assert 0 <= side <= 1

        angle = rotation.normalize_angles(np.array(angle))

        if np.abs(angle) < 1e-4:
            # No need to do anything, the angle is too small to care
            return

        side = side * 2 - 1

        qpos_copy = self.sim.data.qpos.copy()

        # For each cubelet
        for i in range(27):
            cubelet_meta = self.cubelet_meta_info[i]

            if cubelet_meta["type"] == "cubelet":
                mtx = self._cubelet_rotation_matrix(cubelet_meta, qpos_copy)
                current_coords = mtx @ cubelet_meta["coords"].astype(float)
                is_selected = np.take(current_coords, axis) * side > 0.5

                if is_selected:
                    euler = np.zeros(3)
                    euler[axis] = angle

                    combined_matrix = rotation.euler2mat(euler) @ mtx
                    new_euler = rotation.mat2euler(combined_matrix)

                    self.sim.data.qpos[cubelet_meta["euler_qpos"]] = new_euler
            elif cubelet_meta["type"] == "driver":
                # No transformation matrix really here
                current_coords = cubelet_meta["coords"]
                is_selected = np.take(current_coords, axis) * side > 0.5

                if is_selected:
                    joint_idx = self.joints_qpos_idx[cubelet_meta["driver"]]
                    self.sim.data.qpos[joint_idx] += angle