def next_goal()

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

• 40 lines of code
• 10 McCabe index (conditional complexity)

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

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

        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, axis=-1) < face_threshold

        # Z aligned - is there a cube face looking up within the rotation threshold
        z_aligned = rotation.rot_xyz_aligned(cube_quat, rot_threshold)

        axis_nr, axis_sign = cube_utils.up_axis_with_sign(cube_quat)

        cube_aligned = face_aligned and z_aligned

        # Check if current state already meets goal state.
        if cube_aligned and self._is_goal_met(cube_face, face_threshold):
            # Step forward in goal sequence to get next goal.
            self._step_goal()

        goal_action = self._get_goal_action()

        if cube_aligned:
            # Choose 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 face if the face to rotate for next goal is facing up.
            rotate_face = face_to_shift == goal_action.face_idx
        else:
            rotate_face = False

        if rotate_face:
            self.mujoco_simulation.target_model.rotate_face(
                face_to_shift // 2, face_to_shift % 2, goal_action.face_angle
            )

            goal_quat = cube_utils.align_quat_up(cube_quat)
            goal_face = self.goal_face_state
        else:  # need to flip cube

            # Rotate cube so that goal face is on the top. We currently apply
            # a deterministic transformation here that would get the goal face to the top,
            # which is _not_ the minimal possible orientation change, which may be
            # worth addressing in the future.
            goal_quat = self.goal_quat_for_face[goal_action.face_idx]

            # No need to rotate face, just align them.
            goal_face = rounded_current_face

        goal_quat = rotation.quat_normalize(goal_quat)

        return {
            "cube_pos": cube_pos,
            "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,
        }