def _stablize_goal_objects()

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

• 28 lines of code
• 3 McCabe index (conditional complexity)

    def _stablize_goal_objects(self):
        # read current object position and rotation.
        object_pos = self.mujoco_simulation.get_object_pos(pad=False).copy()
        object_quat = self.mujoco_simulation.get_object_quat(pad=False).copy()

        # set object pos/rot as target pos/rot
        self.mujoco_simulation.set_object_pos(
            self.mujoco_simulation.get_target_pos(pad=False)
        )
        self.mujoco_simulation.set_object_quat(
            self.mujoco_simulation.get_target_quat(pad=False)
        )
        self.mujoco_simulation.forward()

        # Recompute z position of objects based on new bounding box computed from new orientation
        updated = update_object_z_coordinate(
            self.mujoco_simulation.get_object_pos(pad=False),
            # Do not use target_bounding_boxes here because we stabilize target by copying its
            # state to object before stabilizing it.
            self.mujoco_simulation.get_object_bounding_boxes(),
            self.mujoco_simulation.get_table_dimensions(),
        )
        self.mujoco_simulation.set_object_pos(updated)
        self.mujoco_simulation.forward()

        stabilize_objects(self.mujoco_simulation)

        # read and set stabilized goal pos/rot
        self.mujoco_simulation.set_target_pos(
            self.mujoco_simulation.get_object_pos(pad=False)
        )
        self.mujoco_simulation.set_target_rot(
            self.mujoco_simulation.get_object_rot(pad=False)
        )
        self.mujoco_simulation.forward()

        # restore object pos/rot back to original
        self.mujoco_simulation.set_object_pos(object_pos)
        self.mujoco_simulation.set_object_quat(object_quat)
        self.mujoco_simulation.forward()