def next_goal()

in robogym/envs/dactyl/goals/face_free.py [0:0]

• 51 lines of code
• 11 McCabe index (conditional complexity)

    def next_goal(self, random_state, current_state):
        """ Generate a new goal from current cube goal state """
        cube_quat = current_state["cube_quat"]
        cube_face = current_state["cube_face_angle"]

        # Success threshold parameters
        face_threshold = self.success_threshold["cube_face_angle"]
        rot_threshold = self.success_threshold["cube_quat"]

        self.mujoco_simulation.clone_target_from_cube()
        self.mujoco_simulation.align_target_faces()

        rounded_current_face = rotation.round_to_straight_angles(cube_face)

        # Face aligned - are faces in the current cube aligned within the threshold
        current_face_diff = rotation.normalize_angles(cube_face - rounded_current_face)
        face_aligned = np.linalg.norm(current_face_diff) < face_threshold

        # Z aligned - is there a cube face looking up within the rotation threshold
        if len(self.face_geom_names) == 2:
            z_aligned = rotation.rot_z_aligned(cube_quat, rot_threshold)
        else:  # len(self.face_geom_names) == 6
            z_aligned = rotation.rot_xyz_aligned(cube_quat, rot_threshold)

        axis_nr, axis_sign = cube_utils.up_axis_with_sign(cube_quat)

        # Do reorientation - with some probability, just reorient the cube
        do_reorientation = random_state.uniform() < self.p_face_flip

        # Rotate face - should we rotate face or reorient the cube
        rotate_face = face_aligned and z_aligned and not do_reorientation

        if rotate_face:
            # Chose index from the geoms that is highest on the z axis
            face_to_shift = cube_utils.face_up(
                self.mujoco_simulation.sim, self.face_geom_names
            )

            # Rotate given face by a random angle and return both, new rotations and an angle
            goal_face, delta_angle = cube_utils.rotated_face_with_angle(
                cube_face,
                face_to_shift,
                random_state,
                self.round_target_face,
                directions=self.goal_directions,
            )

            if len(self.face_geom_names) == 2:
                self.mujoco_simulation.rotate_target_face(face_to_shift, delta_angle)
            else:
                self.mujoco_simulation.rotate_target_face(
                    face_to_shift // 2, face_to_shift % 2, delta_angle
                )

            goal_quat = cube_utils.align_quat_up(cube_quat)
        else:  # need to flip cube
            # Gaol for face rotations is just aligning them
            goal_face = rounded_current_face

            # Make the goal so that a given face is straight up
            candidates = list(range(len(self.face_geom_names)))
            face_to_shift = random_state.choice(candidates)

            z_quat = cube_utils.uniform_z_aligned_quat(random_state)
            face_up_quat = self.goal_quat_for_face[face_to_shift]
            goal_quat = rotation.quat_mul(z_quat, face_up_quat)

        goal_quat = rotation.quat_normalize(goal_quat)

        return {
            "cube_pos": np.zeros(3),
            "cube_quat": goal_quat,
            "cube_face_angle": goal_face,
            "goal_type": "rotation" if rotate_face else "flip",
            "axis_nr": axis_nr,
            "axis_sign": axis_sign,
        }