text: from env.models.tasks.placement_sampler import UniformRandomPegsSampler
  from env.models.tasks import Task, UniformRandomSampler
  from env.mjcf_utils import new_joint, array_to_string
  
  
  class TableTopTask(Task):
      """
      Creates MJCF model of a tabletop task.
  
      A tabletop task consists of one robot interacting with a variable number of
      objects placed on the tabletop. This class combines the robot, the table
      arena, and the objetcts into a single MJCF model.
      """
  
      def __init__(self, mujoco_arena, mujoco_robot, mujoco_objects, mujoco_equality, rng=None, initializer=None):
          """
          Args:
              mujoco_arena: MJCF model of robot workspace
              mujoco_robot: MJCF model of robot model
              mujoco_objects: a list of MJCF models of physical objects
              initializer: placement sampler to initialize object positions.
          """
          super().__init__()
  
          self.merge_arena(mujoco_arena)
          self.merge_robot(mujoco_robot)
          self.merge_objects(mujoco_objects)
          self.merge_equality(mujoco_equality)
          if initializer is None:
              initializer = UniformRandomPegsSampler(rng=rng, ensure_object_boundary_in_range=True, z_rotation=True)
          else:
              print("Warning !")
              
          self.initializer = initializer
          self.initializer.setup(mujoco_objects, self.table_top_offset, self.table_size)
  
      def merge_robot(self, mujoco_robot):
          """Adds robot model to the MJCF model."""
          self.robot = mujoco_robot
          self.merge(mujoco_robot)
  
      def merge_arena(self, mujoco_arena):
          """Adds arena model to the MJCF model."""
          self.arena = mujoco_arena
          self.table_top_offset = mujoco_arena.table_top_abs
          self.table_size = mujoco_arena.table_full_size
          self.merge(mujoco_arena)
  
      def merge_objects(self, mujoco_objects):
          """Adds physical objects to the MJCF model."""
          self.mujoco_objects = mujoco_objects
          self.objects = []  # xml manifestation
          self.targets = []  # xml manifestation
          self.max_horizontal_radius = 0
  
          for obj_name, obj_mjcf in mujoco_objects.items():
              self.merge_asset(obj_mjcf)
              # Load object
              obj = obj_mjcf.get_collision(name=obj_name, site=True)
              obj.append(new_joint(name=obj_name, type="free", damping="0.0001"))
              self.objects.append(obj)
              self.worldbody.append(obj)
  
              self.max_horizontal_radius = max(
                  self.max_horizontal_radius, obj_mjcf.get_horizontal_radius()
              )
  
      def merge_equality(self, equality):
          """Adds equality constraints  to the MJCF model."""
          for one_equality in equality:
              self.equality.append(one_equality)
  
      def place_objects(self):
          """Places objects randomly until no collisions or max iterations hit."""
          return self.initializer.sample()
  