def next_goal()

in robogym/envs/rearrange/goals/object_state.py [0:0]

• 50 lines of code
• 12 McCabe index (conditional complexity)

    def next_goal(self, random_state: RandomState, current_state: dict) -> dict:
        """
        Set goal position for each object and get goal dict.
        """
        goal_valid, goal_dict = self._update_simulation_for_next_goal(random_state)
        target_pos = goal_dict["obj_pos"]
        target_rot = goal_dict["obj_rot"]

        num_objects = self.mujoco_simulation.num_objects
        target_on_table = not self.mujoco_simulation.check_objects_off_table(
            target_pos[:num_objects]
        ).any()

        # Compute which of the goals is within the target placement area. Pad this to include
        # observations for up to max_num_objects (default to `1.0` for empty slots).
        # If an object is out of the placement, the corresponding position in the mask is 0.0.
        in_placement_area = self.mujoco_simulation.check_objects_in_placement_area(
            target_pos, margin=self.args.mask_margin, soft=self.args.soft_mask
        )
        assert in_placement_area.shape == (target_pos.shape[0],)

        # Create qpos for goal state: based on the current qpos but overwrite the object
        # positions to desired positions.
        num_objects = self.mujoco_simulation.num_objects
        qpos_goal = self.mujoco_simulation.qpos.copy()
        for i in range(num_objects):
            qpos_idx = self.mujoco_simulation.mj_sim.model.get_joint_qpos_addr(
                f"object{i}:joint"
            )[0]
            qpos_goal[qpos_idx: qpos_idx + 3] = target_pos[i].copy()
            qpos_goal[qpos_idx + 3: qpos_idx + 7] = rotation.euler2quat(target_rot[i])

        goal_invalid_reason: Optional[str] = None
        if not goal_valid:
            goal_invalid_reason = "Goal placement is invalid"
        elif not target_on_table:
            goal_invalid_reason = "Some goal objects are off the table."

        goal = {
            "obj_pos": target_pos.copy(),
            "obj_rot": target_rot.copy(),
            "qpos_goal": qpos_goal.copy(),
            "goal_valid": goal_valid and target_on_table,
            "goal_in_placement_area": in_placement_area.all(),
            "goal_objects_in_placement_area": in_placement_area.copy(),
            "goal_invalid_reason": goal_invalid_reason,
        }

        if self.args.rot_dist_type == "icp":
            goal["vertices"] = deepcopy(self.mujoco_simulation.get_target_vertices())
            goal["icp"] = [
                ICP(vertices, error_threshold=self.args.icp_error_threshold)
                for vertices in self.mujoco_simulation.get_target_vertices(
                    # Multiple by 2 because in theory max distance between
                    # a vertices and it's closest neighbor should be ~edge length / 2.
                    subdivide_threshold=self.args.icp_error_threshold
                    * 2
                )
            ]

            if self.args.icp_use_bbox_precheck:
                goal[
                    "bounding_box"
                ] = (
                    self.mujoco_simulation.get_target_bounding_boxes_in_table_coordinates().copy()
                )

        return goal