Prompt:
 "colorful-block-tower-on-cylinder-base". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "colorful-block-tower-on-cylinder-base" in a python code block.
###

Response:
 colorful-block-tower-on-cylinder-base: {'task-name': 'colorful-block-tower-on-cylinder-base', 'task-description': 'Construct a tower using four blocks of different colors (red, blue, green, and yellow) on a placed cylindrical base at the corner of the tabletop. The sequence from bottom to top should be red, blue, green, and yellow.', 'assets-used': ['block/block.urdf', 'corner/corner-template.urdf', 'cylinder/cylinder-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class ColorfulBlockTowerOnCylinderBase(Task):
    """Construct a tower using four blocks of different colors (red, blue, green, and yellow) on a placed cylindrical base at the corner of the tabletop. The sequence from bottom to top should be red, blue, green, and yellow."""

    def __init__(self):
        super().__init__()
        self.max_steps = 10
        self.lang_template = "construct a tower using four blocks of different colors (red, blue, green, and yellow) on a placed cylindrical base at the corner of the tabletop. The sequence from bottom to top should be red, blue, green, and yellow."
        self.task_completed_desc = "done building the tower."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add cylindrical base.
        # x, y, z dimensions for the asset size
        base_size = (0.05, 0.05, 0.05)
        base_urdf = 'cylinder/cylinder-template.urdf'
        base_pose = self.get_random_pose(env, base_size)
        base_id = env.add_object(base_urdf, base_pose, 'fixed')

        # Block colors.
        colors = [utils.COLORS['red'], utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['yellow']]

        # Add blocks.
        # x, y, z dimensions for the asset size
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'

        objs = []
        for i in range(4):
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=colors[i])
            objs.append(block_id)

        # Associate placement locations for goals.
        place_pos = [(0, 0, 0.05), (0, 0, 0.09), (0, 0, 0.13), (0, 0, 0.17)]
        targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]

        # Goal: blocks are stacked on the cylindrical base in the order red, blue, green, yellow from bottom to top.
        for i in range(4):
            self.add_goal(objs=[objs[i]], matches=np.ones((1, 1)), targ_poses=[targs[i]], replace=False,
                    rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2])
            self.lang_goals.append(self.lang_template)```
 END


Prompt:
 "align-cylinders-in-square". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "align-cylinders-in-square" in a python code block.
###

Response:
 align-cylinders-in-square: {'task-name': 'align-cylinders-in-square', 'task-description': 'Arrange four cylinders of different colors (red, blue, green, yellow) on the corners of a square facing the center. The square is outlined by four zones of matching colors at each corner. The red cylinder should be placed at the red zone and facing the center of the square, and same for the remaining colors.', 'assets-used': ['cylinder/cylinder-template.urdf', 'zone/zone.urdf', 'square/square-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class AlignCylindersInSquare(Task):
    """Arrange four cylinders of different colors (red, blue, green, yellow)
    on the corners of a square facing the center. The square is outlined by
    four zones of matching colors at each corner. The red cylinder should be
    placed at the red zone and facing the center of the square, and same for
    the remaining colors."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "arrange the {color} cylinder at the {color} corner of the square"
        self.task_completed_desc = "done arranging cylinders."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors and corresponding zones
        colors = ['red', 'blue', 'green', 'yellow']
        zones = ['zone-red', 'zone-blue', 'zone-green', 'zone-yellow']

        # Add zones at the corners of a square
        zone_size = (0.05, 0.05, 0.005)
        zone_urdf = 'zone/zone.urdf'
        zone_poses = []
        for i in range(4):
            zone_pose = self.get_random_pose(env, zone_size)
            env.add_object(zone_urdf, zone_pose, 'fixed')
            zone_poses.append(zone_pose)

        # Add cylinders of different colors
        cylinder_size = (0.02, 0.02, 0.08)
        cylinder_urdf = 'cylinder/cylinder-template.urdf'
        cylinders = []
        for i in range(4):
            cylinder_pose = self.get_random_pose(env, cylinder_size)
            color = utils.COLORS[colors[i]]
            cylinder_id = env.add_object(cylinder_urdf, cylinder_pose, color=color)
            cylinders.append(cylinder_id)

        # Goal: each cylinder is in a different zone of the same color
        for i in range(4):
            self.add_goal(objs=[cylinders[i]], matches=np.ones((1, 1)), targ_poses=[zone_poses[i]], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1/4)
            self.lang_goals.append(self.lang_template.format(color=colors[i]))```
 END


Prompt:
 "construct-corner-blocks". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "construct-corner-blocks" in a python code block.
###

Response:
 construct-corner-blocks: {'task-name': 'construct-corner-blocks', 'task-description': 'Create a corner structure using four blocks. Two red blocks form the base, one on each side of the corner, followed by a green block that is positioned on the red blocks at the corner junction, and finally a blue block on top of the green one. The overall structure forms a 3-D corner.', 'assets-used': ['block/block.urdf', 'corner/corner-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class ConstructCornerBlocks(Task):
    """Create a corner structure using four blocks. Two red blocks form the base, one on each side of the corner, followed by a green block that is positioned on the red blocks at the corner junction, and finally a blue block on top of the green one. The overall structure forms a 3-D corner."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "Create a corner structure using four blocks. Two red blocks form the base, one on each side of the corner, followed by a green block that is positioned on the red blocks at the corner junction, and finally a blue block on top of the green one. The overall structure forms a 3-D corner."
        self.task_completed_desc = "done constructing corner blocks."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add corner.
        corner_size = (0.15, 0.15, 0.05)
        corner_urdf = 'corner/corner-template.urdf'
        corner_pose = self.get_random_pose(env, corner_size)
        env.add_object(corner_urdf, corner_pose, 'fixed')

        # Add blocks.
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        block_colors = [utils.COLORS['red'], utils.COLORS['red'], utils.COLORS['green'], utils.COLORS['blue']]
        blocks = []
        for i in range(4):
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=block_colors[i])
            blocks.append(block_id)

        # Associate placement locations for goals.
        place_pos = [(0, -0.05, 0.02), (0, 0.05, 0.02), (0, 0, 0.06), (0, 0, 0.10)]
        targs = [(utils.apply(corner_pose, i), corner_pose[1]) for i in place_pos]

        # Goal: blocks are stacked in a corner (bottom row: two red blocks).
        self.add_goal(objs=blocks[:2], matches=np.ones((2, 2)), targ_poses=targs[:2], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 2, symmetries=[np.pi/2]*2)

        # Goal: blocks are stacked in a corner (middle row: one green block).
        self.add_goal(objs=blocks[2:3], matches=np.ones((1, 1)), targ_poses=targs[2:3], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 3, symmetries=[np.pi/2]*1)

        # Goal: blocks are stacked in a corner (top row: one blue block).
        self.add_goal(objs=blocks[3:], matches=np.ones((1, 1)), targ_poses=targs[3:], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2]*1)```
 END


Prompt:
 "corner-sort-cylinders". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "corner-sort-cylinders" in a python code block.
###

Response:
 corner-sort-cylinders: {'task-name': 'corner-sort-cylinders', 'task-description': 'Pick up cylinders of four different colors (red, blue, green, yellow) and place them into four corners accordingly marked on the tabletop. The corner is designed as a 2-block-size square where only one cylinder can fit. The task is challenging due to precise placement and the need for identifying the corners accurately.', 'assets-used': ['corner/corner-template.urdf', 'cylinder/cylinder-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class CornerSortCylinders(Task):
    """Pick up cylinders of four different colors (red, blue, green, yellow) and place them into four corners accordingly marked on the tabletop."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "place the {color} cylinder in the {color} corner"
        self.task_completed_desc = "done sorting cylinders."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors
        colors = ['red', 'blue', 'green', 'yellow']

        # Add corners
        corner_size = (0.04, 0.04, 0.04)  # x, y, z dimensions for the asset size
        corner_template = 'corner/corner-template.urdf'
        corner_poses = []
        for color in colors:
            replace = {'DIM': corner_size, 'HALF': (corner_size[0] / 2, corner_size[1] / 2, corner_size[2] / 2), 'COLOR': utils.COLORS[color]}
            corner_urdf = self.fill_template(corner_template, replace)
            corner_pose = self.get_random_pose(env, corner_size)
            env.add_object(corner_urdf, corner_pose, 'fixed')
            corner_poses.append(corner_pose)

        # Add cylinders
        cylinder_size = (0.02, 0.02, 0.06)  # x, y, z dimensions for the asset size
        cylinder_template = 'cylinder/cylinder-template.urdf'
        cylinders = []
        for color in colors:
            replace = {'DIM': cylinder_size, 'HALF': (cylinder_size[0] / 2, cylinder_size[1] / 2, cylinder_size[2] / 2), 'COLOR': utils.COLORS[color]}
            cylinder_urdf = self.fill_template(cylinder_template, replace)
            cylinder_pose = self.get_random_pose(env, cylinder_size)
            cylinder_id = env.add_object(cylinder_urdf, cylinder_pose)
            cylinders.append(cylinder_id)

        # Add goals
        for i in range(len(cylinders)):
            self.add_goal(objs=[cylinders[i]], matches=np.int32([[1]]), targ_poses=[corner_poses[i]], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1 / len(cylinders))
            self.lang_goals.append(self.lang_template.format(color=colors[i]))```
 END


Prompt:
 "sorting-blocks-into-pallets". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "sorting-blocks-into-pallets" in a python code block.
###

Response:
 sorting-blocks-into-pallets: {'task-name': 'sorting-blocks-into-pallets', 'task-description': 'Pick up blocks of four different colors (red, blue, green, yellow) and place them into four separate pallets of matching color. The pallets are placed in a row and the blocks are scattered randomly on the table.', 'assets-used': ['block/block.urdf', 'pallet/pallet.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class SortingBlocksIntoPallets(Task):
    """Pick up blocks of four different colors (red, blue, green, yellow) and place them into four separate pallets of matching color. The pallets are placed in a row and the blocks are scattered randomly on the table."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "put the {color} block into the {color} pallet"
        self.task_completed_desc = "done sorting blocks into pallets."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add pallets.
        # x, y, z dimensions for the asset size
        pallet_size = (0.12, 0.12, 0.02)
        pallet_urdf = 'pallet/pallet.urdf'
        pallet_poses = []
        pallet_colors = ['red', 'blue', 'green', 'yellow']
        for color in pallet_colors:
            pallet_pose = self.get_random_pose(env, pallet_size)
            env.add_object(pallet_urdf, pallet_pose, 'fixed', color=utils.COLORS[color])
            pallet_poses.append(pallet_pose)

        # Add blocks.
        # x, y, z dimensions for the asset size
        blocks = []
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        for color in pallet_colors:
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=utils.COLORS[color])
            blocks.append(block_id)

        # Goal: each block is in a different pallet of matching color.
        for i in range(len(blocks)):
            self.add_goal(objs=[blocks[i]], matches=np.ones((1, 1)), targ_poses=[pallet_poses[i]], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1/len(blocks))
            self.lang_goals.append(self.lang_template.format(color=pallet_colors[i]))```
 END


Prompt:
 "sort-and-assemble-block-castle". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "sort-and-assemble-block-castle" in a python code block.
###

Response:
 sort-and-assemble-block-castle: {'task-name': 'sort-and-assemble-block-castle', 'task-description': 'On a tabletop, there are twelve blocks of three different colors (four red, four green, and four blue). The task involves sorting the blocks according to the color in three marked zones on the tabletop and subsequently constructing a castle in each zone. In each castle, the first layer should consist of the two blocks of the same color, followed by the second layer of one block and finally the last block on the top forming a castle-like structure. The challenge lies in the color-based sorting and careful assembly of the blocks to avoid toppling.', 'assets-used': ['block/block.urdf', 'zone/zone.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class SortAndAssembleBlockCastle(Task):
    """Sort blocks by color and assemble them into a castle-like structure."""

    def __init__(self):
        super().__init__()
        self.max_steps = 50
        self.lang_template = "sort the blocks by color and assemble them into a castle"
        self.task_completed_desc = "done sorting and assembling."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add zones.
        zone_size = (0.12, 0.12, 0)
        zone_urdf = 'zone/zone.urdf'
        zone_poses = []
        for _ in range(3):
            zone_pose = self.get_random_pose(env, zone_size)
            env.add_object(zone_urdf, zone_pose, 'fixed')
            zone_poses.append(zone_pose)

        # Add blocks.
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        block_colors = [utils.COLORS['red'], utils.COLORS['green'], utils.COLORS['blue']]
        blocks = []
        for color in block_colors:
            for _ in range(4):
                block_pose = self.get_random_pose(env, block_size)
                block_id = env.add_object(block_urdf, block_pose, color=color)
                blocks.append(block_id)

        # Goal: each block is in a different zone based on color.
        for i in range(12):
            self.add_goal(objs=[blocks[i]], matches=np.ones((1, 3)), targ_poses=zone_poses, replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1/12)

        # Goal: blocks are stacked in a pyramid in each zone.
        for i in range(3):
            zone_blocks = blocks[i*4:(i+1)*4]
            place_pos = [(0, -0.02, 0.02), (0, 0.02, 0.02),
                         (0, 0, 0.06), (0, 0, 0.10)]
            targs = [(utils.apply(zone_poses[i], pos), zone_poses[i][1]) for pos in place_pos]
            for j in range(4):
                self.add_goal(objs=[zone_blocks[j]], matches=np.ones((1, 1)), targ_poses=[targs[j]], replace=False,
                              rotations=True, metric='pose', params=None, step_max_reward=1/12)

        self.lang_goals.append(self.lang_template)```
 END


Prompt:
 "vertical-insertion-blocks". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "vertical-insertion-blocks" in a python code block.
###

Response:
 vertical-insertion-blocks: {'task-name': 'vertical-insertion-blocks', 'task-description': 'The task involves picking up four color specific blocks (red, blue, green, and yellow) and inserting each block into four differently colored stands set upright on the tabletop. The block-colored with red needs to be inserted into the red-colored stand, and the same sequence is maintained for each colored blocks and stands. This task is challenging due to the requirement for precise insertion and the manipulation of vertical objects.', 'assets-used': ['stacking/block.urdf', 'stacking/stand.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
import pybullet as p

class VerticalInsertionBlocks(Task):
    """Pick up four color specific blocks and insert each block into four differently colored stands set upright on the tabletop."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "insert the {color} block into the {color} stand"
        self.task_completed_desc = "done inserting blocks into stands."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors for blocks and stands
        colors = ['red', 'blue', 'green', 'yellow']

        # Add stands.
        # x, y, z dimensions for the asset size
        stand_size = (0.04, 0.04, 0.1)
        stand_urdf = 'stacking/stand.urdf'
        stands = []
        for color in colors:
            stand_pose = self.get_random_pose(env, stand_size)
            stand_id = env.add_object(stand_urdf, stand_pose, color=utils.COLORS[color], category='fixed')
            stands.append(stand_id)

        # Add blocks.
        # x, y, z dimensions for the asset size
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'stacking/block.urdf'
        blocks = []
        for color in colors:
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=utils.COLORS[color])
            blocks.append(block_id)

        # Goal: each block is inserted into the stand of the same color.
        for i in range(len(blocks)):
            self.add_goal(objs=[blocks[i]], matches=np.ones((1, 1)), targ_poses=[p.getBasePositionAndOrientation(stands[i])], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1/len(blocks))
            self.lang_goals.append(self.lang_template.format(color=colors[i]))```
 END


Prompt:
 "color-coordinated-sphere-insertion". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "color-coordinated-sphere-insertion" in a python code block.
###

Response:
 color-coordinated-sphere-insertion: {'task-name': 'color-coordinated-sphere-insertion', 'task-description': 'There are four spheres and four boxes of different colors (red, blue, green, and yellow). Each sphere is inside a box, but not corresponding to the color of the box. The task is to pick up each sphere and place it in the box of the same color. The challenge lies in the precise placement and color matching.', 'assets-used': ['sphere/sphere.urdf', 'box/box-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class ColorCoordinatedSphereInsertion(Task):
    """Insert each sphere into the bowl of the same color."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "insert each sphere into the bowl of the same color"
        self.task_completed_desc = "done inserting spheres into bowls."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors and corresponding names
        colors = ['red', 'blue', 'green', 'yellow']
        color_values = [utils.COLORS[color] for color in colors]

        # Add bowls.
        # x, y, z dimensions for the asset size
        bowl_size = (0.12, 0.12, 0.02)
        bowl_urdf = 'bowl/bowl.urdf'
        bowl_poses = []
        for i in range(4):
            bowl_pose = self.get_random_pose(env, bowl_size)
            env.add_object(bowl_urdf, bowl_pose, 'fixed', color=color_values[i])
            bowl_poses.append(bowl_pose)

        # Add spheres.
        # x, y, z dimensions for the asset size
        sphere_size = (0.04, 0.04, 0.04)
        sphere_template = 'sphere/sphere-template.urdf'
        spheres = []
        for i in range(4):
            sphere_pose = self.get_random_pose(env, sphere_size)
            replace = {'DIM': sphere_size, 'HALF': (sphere_size[0] / 2, sphere_size[1] / 2, sphere_size[2] / 2)}
            sphere_urdf = self.fill_template(sphere_template, replace)
            sphere_id = env.add_object(sphere_urdf, sphere_pose, color=color_values[i])
            spheres.append(sphere_id)

        # Goal: each sphere is in a bowl of the same color.
        for i in range(4):
            self.add_goal(objs=[spheres[i]], matches=np.ones((1, 1)), targ_poses=[bowl_poses[i]], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1/4)
            self.lang_goals.append(f"insert the {colors[i]} sphere into the {colors[i]} bowl")```
 END


Prompt:
 "block-pyramid-with-limited-space". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "block-pyramid-with-limited-space" in a python code block.
###

Response:
 block-pyramid-with-limited-space: {'task-name': 'block-pyramid-with-limited-space', 'task-description': 'On a tabletop, there are twelve blocks of four different colors (three red, three green, three blue, three yellow). Three zones are defined, each with a triangle-shaped border that is marked. The task involves sorting the blocks according to the color into three zones on the tabletop and constructing a pyramid in each zone. In each pyramid, the base should contain two blocks of the same color, followed by the second layer of one block, thus forming a pyramid-like structure. However, the third yellow block should be placed in the center of the zones forming a smaller pyramid. The challenge lies in the color-based sorting, careful assembly of the blocks to avoid topple, and limited space in the zones which adds to the complexity.', 'assets-used': ['block/block.urdf', 'zone/zone.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class BlockPyramidWithLimitedSpace(Task):
    """Sort blocks according to color into three zones on the tabletop and construct a pyramid in each zone."""

    def __init__(self):
        super().__init__()
        self.max_steps = 50
        self.lang_template = "sort the blocks according to color into three zones and construct a pyramid in each zone"
        self.task_completed_desc = "done sorting and constructing pyramids."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add zones.
        zone_size = (0.12, 0.12, 0)
        zone_urdf = 'zone/zone.urdf'
        zone_poses = []
        for _ in range(3):
            zone_pose = self.get_random_pose(env, zone_size)
            env.add_object(zone_urdf, zone_pose, 'fixed')
            zone_poses.append(zone_pose)

        # Block colors.
        colors = [
            utils.COLORS['red'], utils.COLORS['green'], utils.COLORS['blue'], utils.COLORS['yellow']
        ]

        # Add blocks.
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        blocks = []
        for color in colors:
            for _ in range(3):
                block_pose = self.get_random_pose(env, block_size)
                block_id = env.add_object(block_urdf, block_pose, color=color)
                blocks.append(block_id)

        # Associate placement locations for goals.
        place_pos = [(0, -0.05, 0.03), (0, 0, 0.03),
                     (0, 0.05, 0.03), (0, -0.025, 0.08),
                     (0, 0.025, 0.08), (0, 0, 0.13)]
        targs = [(utils.apply(zone_pose, i), zone_pose[1]) for zone_pose in zone_poses for i in place_pos]

        # Goal: blocks are sorted and stacked in a pyramid in each zone.
        for i in range(3):
            self.add_goal(objs=blocks[i*3:(i+1)*3], matches=np.ones((3, 3)), targ_poses=targs[i*3:(i+1)*3], replace=False,
                    rotations=True, metric='pose', params=None, step_max_reward=1 / 3, symmetries=[np.pi/2]*3)
        self.lang_goals.append(self.lang_template)```
 END


Prompt:
 "build-cylinder-structure". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "build-cylinder-structure" in a python code block.
###

Response:
 build-cylinder-structure: {'task-name': 'build-cylinder-structure', 'task-description': "Using four colored cylinders (red, blue, green, yellow), construct a structure atop a square base. The red and blue cylinders should be sealed by the square base side by side, while the green cylinder should be on top of the blue one, and the yellow one on top of the red. The final look should resemble the letter 'H'.", 'assets-used': ['cylinder/cylinder-template.urdf', 'square/square-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class BuildCylinderStructure(Task):
    """Construct a structure using four colored cylinders (red, blue, green, yellow) on a square base."""

    def __init__(self):
        super().__init__()
        self.max_steps = 5
        self.lang_template = "construct a structure using four colored cylinders on a square base"
        self.task_completed_desc = "done building the cylinder structure."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add square base.
        # x, y, z dimensions for the asset size
        base_size = (0.15, 0.15, 0.005)
        base_urdf = 'square/square-template.urdf'
        base_pose = self.get_random_pose(env, base_size)
        env.add_object(base_urdf, base_pose, category='fixed')

        # Cylinder colors.
        colors = [
            utils.COLORS['red'], utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['yellow']
        ]

        # Add cylinders.
        # x, y, z dimensions for the asset size
        cylinder_size = (0.04, 0.04, 0.08)
        cylinder_urdf = 'cylinder/cylinder-template.urdf'

        objs = []
        for i in range(4):
            cylinder_pose = self.get_random_pose(env, cylinder_size)
            cylinder_id = env.add_object(cylinder_urdf, cylinder_pose, color=colors[i])
            objs.append(cylinder_id)

        # Associate placement locations for goals.
        place_pos = [(0, -0.05, 0.04), (0, 0.05, 0.04),
                     (0, 0.05, 0.12), (0, -0.05, 0.12)]
        targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]

        # Goal: red and blue cylinders are placed side by side on the base.
        self.add_goal(objs=objs[:2], matches=np.ones((2, 2)), targ_poses=targs[:2], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 2, symmetries=[np.pi/2]*2)
        self.lang_goals.append("place the red and blue cylinders side by side on the base")

        # Goal: green cylinder is placed on top of the blue cylinder.
        self.add_goal(objs=[objs[2]], matches=np.ones((1, 1)), targ_poses=[targs[2]], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 2, symmetries=[np.pi/2])
        self.lang_goals.append("place the green cylinder on top of the blue cylinder")

        # Goal: yellow cylinder is placed on top of the red cylinder.
        self.add_goal(objs=[objs[3]], matches=np.ones((1, 1)), targ_poses=[targs[3]], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 2, symmetries=[np.pi/2])
        self.lang_goals.append("place the yellow cylinder on top of the red cylinder")```
 END


Prompt:
 "insert-blocks-lineup". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "insert-blocks-lineup" in a python code block.
###

Response:
 insert-blocks-lineup: {'task-name': 'insert-blocks-lineup', 'task-description': 'On the tabletop, there are four different color blocks (red, blue, green, and yellow), and four fixtures in corresponding colors. The task is to pick up each block and insert it into the corresponding color fixture. However, the fixtures are arranged in a straight line, with a line of small blocks serving as barrier between the fixtures and the colored blocks initially scattered on the table, providing a challenge in precise navigation and placement.', 'assets-used': ['block/block.urdf', 'insertion/fixture.urdf', 'block/small.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class InsertBlocksLineup(Task):
    """Pick up four different color blocks and insert them into the corresponding color fixtures."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "insert the {color} block into the {color} fixture"
        self.task_completed_desc = "done inserting blocks."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors for blocks and fixtures
        colors = ['red', 'blue', 'green', 'yellow']

        # Add fixtures.
        fixture_size = (0.04, 0.04, 0.04)
        fixture_urdf = 'insertion/fixture.urdf'
        fixture_poses = []
        for i in range(4):
            fixture_pose = self.get_random_pose(env, fixture_size)
            fixture_id = env.add_object(fixture_urdf, fixture_pose, color=utils.COLORS[colors[i]], category='fixed')
            fixture_poses.append((fixture_pose, fixture_id))

        # Add blocks.
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        blocks = []
        for i in range(4):
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=utils.COLORS[colors[i]])
            blocks.append(block_id)

        # Add small blocks as barriers.
        small_block_size = (0.02, 0.02, 0.02)
        small_block_urdf = 'block/small.urdf'
        for _ in range(10):
            small_block_pose = self.get_random_pose(env, small_block_size)
            env.add_object(small_block_urdf, small_block_pose)

        # Goal: each block is in the corresponding color fixture.
        for i in range(4):
            self.add_goal(objs=[blocks[i]], matches=np.ones((1, 1)), targ_poses=[fixture_poses[i][0]], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1/4)
            self.lang_goals.append(self.lang_template.format(color=colors[i]))```
 END


Prompt:
 "color-specific-container-fill". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "color-specific-container-fill" in a python code block.
###

Response:
 color-specific-container-fill: {'task-name': 'color-specific-container-fill', 'task-description': 'Arrange four colored blocks (red, blue, green, and yellow) around a pallet. Then, pick up these blocks and place them inside a container marked in the same color. The task requires precise placement, color matching, and an understanding of spatial structures.', 'assets-used': ['block/block.urdf', 'pallet/pallet.urdf', 'container/container-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class ColorSpecificContainerFill(Task):
    """Arrange four colored blocks (red, blue, green, and yellow) around a pallet.
    Then, pick up these blocks and place them inside a container marked in the same color.
    The task requires precise placement, color matching, and an understanding of spatial structures."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "put the {color} block in the {color} container"
        self.task_completed_desc = "done arranging blocks in containers."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add pallet.
        # x, y, z dimensions for the asset size
        pallet_size = (0.15, 0.15, 0.01)
        pallet_urdf = 'pallet/pallet.urdf'
        pallet_pose = self.get_random_pose(env, pallet_size)
        env.add_object(pallet_urdf, pallet_pose, 'fixed')

        # Define block and container colors.
        colors = ['red', 'blue', 'green', 'yellow']

        # Add blocks and containers.
        # x, y, z dimensions for the asset size
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        container_size = (0.12, 0.12, 0.05)
        container_template = 'container/container-template.urdf'
        blocks = []
        containers = []
        for color in colors:
            # Add block.
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=utils.COLORS[color])
            blocks.append(block_id)

            # Add container.
            container_pose = self.get_random_pose(env, container_size)
            replace = {'DIM': container_size, 'HALF': (container_size[0] / 2, container_size[1] / 2, container_size[2] / 2)}
            container_urdf = self.fill_template(container_template, replace)
            container_id = env.add_object(container_urdf, container_pose, 'fixed', color=utils.COLORS[color])
            containers.append(container_id)

        # Goal: each block is in a container of the same color.
        for i in range(len(colors)):
            self.add_goal(objs=[blocks[i]], matches=np.ones((1, 1)), targ_poses=[p.getBasePositionAndOrientation(containers[i])], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / len(colors))
            self.lang_goals.append(self.lang_template.format(color=colors[i]))```
 END


Prompt:
 "multicolor-block-bridge". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "multicolor-block-bridge" in a python code block.
###

Response:
 multicolor-block-bridge: {'task-name': 'multicolor-block-bridge', 'task-description': 'Build a bridge by stacking three red, three blue, and three green blocks on a pallet. Arrange in a sequence from left to right: red, blue, and green. Then, place three cylinders of corresponding colors on top of the stacked blocks, forming a bridge. The cylinders should roll from the top block to the pallet, creating a challenge of precision and control.', 'assets-used': ['block/block.urdf', 'pallet/pallet.urdf', 'cylinder/cylinder-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class MulticolorBlockBridge(Task):
    """Build a bridge by stacking three red, three blue, and three green blocks on a pallet.
    Arrange in a sequence from left to right: red, blue, and green.
    Then, place three cylinders of corresponding colors on top of the stacked blocks, forming a bridge.
    The cylinders should roll from the top block to the pallet, creating a challenge of precision and control."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "Build a bridge by stacking three red, three blue, and three green blocks on a pallet. Arrange in a sequence from left to right: red, blue, and green. Then, place three cylinders of corresponding colors on top of the stacked blocks, forming a bridge."
        self.task_completed_desc = "done building the bridge."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add pallet.
        # x, y, z dimensions for the asset size
        pallet_size = (0.15, 0.15, 0.01)
        pallet_urdf = 'pallet/pallet.urdf'
        pallet_pose = self.get_random_pose(env, pallet_size)
        env.add_object(pallet_urdf, pallet_pose, 'fixed')

        # Add blocks.
        # x, y, z dimensions for the asset size
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        block_colors = [utils.COLORS['red'], utils.COLORS['blue'], utils.COLORS['green']]
        blocks = []
        for i in range(9):  # 3 blocks of each color
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=block_colors[i // 3])
            blocks.append(block_id)

        # Add cylinders.
        # x, y, z dimensions for the asset size
        cylinder_size = (0.04, 0.04, 0.04)
        cylinder_template = 'cylinder/cylinder-template.urdf'
        cylinders = []
        for i in range(3):  # 1 cylinder of each color
            cylinder_pose = self.get_random_pose(env, cylinder_size)
            replace = {'DIM': cylinder_size, 'HALF': (cylinder_size[0] / 2, cylinder_size[1] / 2, cylinder_size[2] / 2)}
            cylinder_urdf = self.fill_template(cylinder_template, replace)
            cylinder_id = env.add_object(cylinder_urdf, cylinder_pose, color=block_colors[i])
            cylinders.append(cylinder_id)

        # Associate placement locations for goals.
        place_pos = [(0, -0.05, 0.03), (0, 0, 0.03), (0, 0.05, 0.03)]
        targs = [(utils.apply(pallet_pose, i), pallet_pose[1]) for i in place_pos]

        # Goal: blocks are stacked on the pallet in the order red, blue, green.
        for i in range(9):
            self.add_goal(objs=[blocks[i]], matches=np.ones((1, 1)), targ_poses=[targs[i // 3]], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1 / 9, symmetries=[np.pi/2])

        # Goal: cylinders are placed on top of the stacked blocks.
        for i in range(3):
            self.add_goal(objs=[cylinders[i]], matches=np.ones((1, 1)), targ_poses=[targs[i]], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1 / 3, symmetries=[np.pi/2])

        self.lang_goals.append(self.lang_template)```
 END


Prompt:
 "pyramid-blocks-assemble". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "pyramid-blocks-assemble" in a python code block.
###

Response:
 pyramid-blocks-assemble: {'task-name': 'pyramid-blocks-assemble', 'task-description': 'Construct a pyramid using nine blocks in a specific color order on a pallet. The bottom layer should contain five blocks: red, blue, green, yellow, and orange (in that order from left to right). The middle layer should consist of three blocks: yellow, red, and blue (from left to right). The top layer should contain a single green block. The pyramid requires careful placement and color matching.', 'assets-used': ['block/block.urdf', 'pallet/pallet.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class PyramidBlocksAssemble(Task):
    """Construct a pyramid using nine blocks in a specific color order on a pallet."""

    def __init__(self):
        super().__init__()
        self.max_steps = 15
        self.lang_template = "construct a pyramid using nine blocks in a specific color order on a pallet"
        self.task_completed_desc = "done constructing pyramid."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add pallet.
        # x, y, z dimensions for the asset size
        pallet_size = (0.35, 0.35, 0.01)
        pallet_urdf = 'pallet/pallet.urdf'
        pallet_pose = self.get_random_pose(env, pallet_size)
        env.add_object(pallet_urdf, pallet_pose, 'fixed')

        # Block colors.
        colors = [
            utils.COLORS['red'], utils.COLORS['blue'], utils.COLORS['green'],
            utils.COLORS['yellow'], utils.COLORS['orange']
        ]

        # Add blocks.
        # x, y, z dimensions for the asset size
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        blocks = []
        for _ in range(9):
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose)
            blocks.append(block_id)

        # Associate placement locations for goals.
        place_pos = [
            (-0.1, -0.1, 0.02), (0, -0.1, 0.02), (0.1, -0.1, 0.02), (-0.1, 0, 0.02), (0.1, 0, 0.02),
            (-0.05, 0.05, 0.06), (0.05, 0.05, 0.06), (0, 0.1, 0.06),
            (0, 0.05, 0.1)
        ]
        targs = [(utils.apply(pallet_pose, i), pallet_pose[1]) for i in place_pos]

        # Goal: blocks are stacked in a pyramid in a specific color order.
        for i in range(9):
            self.add_goal(objs=[blocks[i]], matches=np.ones((1, 1)), targ_poses=[targs[i]], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1 / 9)
            self.lang_goals.append(self.lang_template.format(blocks="the blocks",
                                                             row="row"))```
 END


Prompt:
 "triangle-block-arrangement". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "triangle-block-arrangement" in a python code block.
###

Response:
 triangle-block-arrangement: {'task-name': 'triangle-block-arrangement', 'task-description': 'Arrange blocks of three different colors (red, green, and blue) in a triangular layout on the tabletop. The base of the triangle should consist of three red blocks, followed by a row of two green blocks, and finally a single blue block at the top of the triangle. The task is challenging due to the need for precise placement and color coordination.', 'assets-used': ['block/block.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class TriangleBlockArrangement(Task):
    """Arrange blocks of three different colors (red, green, and blue) in a triangular layout on the tabletop."""

    def __init__(self):
        super().__init__()
        self.max_steps = 6
        self.lang_template = "Arrange blocks of three different colors (red, green, and blue) in a triangular layout on the tabletop."
        self.task_completed_desc = "done arranging blocks in a triangle."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add blocks.
        # x, y, z dimensions for the asset size
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        colors = [utils.COLORS['red'], utils.COLORS['green'], utils.COLORS['blue']]
        blocks = []
        for i in range(6):
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=colors[i//2])
            blocks.append(block_id)

        # Associate placement locations for goals.
        place_pos = [(0, -0.05, 0.02), (0, 0, 0.02),
                     (0, 0.05, 0.02), (0, -0.025, 0.06),
                     (0, 0.025, 0.06), (0, 0, 0.10)]
        base_pose = self.get_random_pose(env, block_size)
        targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]

        # Goal: blocks are stacked in a triangle (bottom row: red, middle row: green, top row: blue).
        self.add_goal(objs=blocks[:3], matches=np.ones((3, 3)), targ_poses=targs[:3], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 2, symmetries=[np.pi/2]*3)
        self.lang_goals.append("Arrange the red blocks in a row at the base of the triangle.")

        self.add_goal(objs=blocks[3:5], matches=np.ones((2, 2)), targ_poses=targs[3:5], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 3, symmetries=[np.pi/2]*2)
        self.lang_goals.append("Arrange the green blocks in a row above the red blocks.")

        self.add_goal(objs=blocks[5:], matches=np.ones((1, 1)), targ_poses=targs[5:], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 6, symmetries=[np.pi/2])
        self.lang_goals.append("Place the blue block at the top of the triangle.")```
 END


Prompt:
 "place-ball-in-elevated-bowl". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "place-ball-in-elevated-bowl" in a python code block.
###

Response:
 place-ball-in-elevated-bowl: {'task-name': 'place-ball-in-elevated-bowl', 'task-description': 'The primary objective of the task is to pick up a red ball and carefully place it into a bowl, which is positioned on a raised platform that is surrounded by small blocks. The challenge lies in precise navigation, maintaining a hold of the ball, and avoiding the surrounding blocks.', 'assets-used': ['ball/ball-template.urdf', 'bowl/bowl.urdf', 'block/small.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
import pybullet as p

class PlaceBallInElevatedBowl(Task):
    """Pick up a red ball and carefully place it into a bowl, which is positioned on a raised platform that is surrounded by small blocks."""

    def __init__(self):
        super().__init__()
        self.max_steps = 10
        self.lang_template = "place the red ball in the elevated bowl"
        self.task_completed_desc = "done placing ball in bowl."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add elevated platform.
        platform_size = (0.3, 0.3, 0.05)
        platform_urdf = 'platform/platform.urdf'
        platform_pose = self.get_random_pose(env, platform_size)
        env.add_object(platform_urdf, platform_pose, 'fixed')

        # Add bowl on the platform.
        bowl_size = (0.12, 0.12, 0)
        bowl_urdf = 'bowl/bowl.urdf'
        bowl_pose = self.get_random_pose(env, bowl_size)
        bowl_pose[0][2] += platform_size[2]  # place the bowl on top of the platform
        bowl_id = env.add_object(bowl_urdf, bowl_pose, 'fixed')

        # Add red ball.
        ball_size = (0.04, 0.04, 0.04)
        ball_urdf = 'ball/ball-template.urdf'
        ball_pose = self.get_random_pose(env, ball_size)
        ball_id = env.add_object(ball_urdf, ball_pose, color=utils.COLORS['red'])

        # Add small blocks around the platform.
        block_size = (0.02, 0.02, 0.02)
        block_urdf = 'block/small.urdf'
        for _ in range(5):
            block_pose = self.get_random_pose(env, block_size)
            env.add_object(block_urdf, block_pose)

        # Goal: the red ball is in the bowl.
        self.add_goal(objs=[ball_id], matches=np.ones((1, 1)), targ_poses=[bowl_pose], replace=False,
                      rotations=True, metric='pose', params=None, step_max_reward=1)
        self.lang_goals.append(self.lang_template)```
 END


Prompt:
 "align-balls-in-colored-zones". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "align-balls-in-colored-zones" in a python code block.
###

Response:
 align-balls-in-colored-zones: {'task-name': 'align-balls-in-colored-zones', 'task-description': 'There are six balls of different colors (red, blue, green, yellow, orange, and purple) and six zones correspondingly colored. The task is to pick up each ball and place it in the zone of the same color, arranging the balls in a straight line. The challenge lies in the precise placement and color matching.', 'assets-used': ['ball/ball-template.urdf', 'zone/zone.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class AlignBallsInColoredZones(Task):
    """Align balls of different colors in correspondingly colored zones."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "place the {color} ball in the {color} zone"
        self.task_completed_desc = "done aligning balls in colored zones."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors for balls and zones
        colors = ['red', 'blue', 'green', 'yellow', 'orange', 'purple']
        color_names = ['Red', 'Blue', 'Green', 'Yellow', 'Orange', 'Purple']

        # Add zones.
        zone_size = (0.12, 0.12, 0)
        zone_urdf = 'zone/zone.urdf'
        zone_poses = []
        for i in range(6):
            zone_pose = self.get_random_pose(env, zone_size)
            env.add_object(zone_urdf, zone_pose, 'fixed', color=utils.COLORS[colors[i]])
            zone_poses.append(zone_pose)

        # Add balls.
        ball_size = (0.04, 0.04, 0.04)
        ball_urdf = 'ball/ball-template.urdf'
        balls = []
        for i in range(6):
            ball_pose = self.get_random_pose(env, ball_size)
            replace = {'DIM': ball_size, 'HALF': (ball_size[0] / 2, ball_size[1] / 2, ball_size[2] / 2), 'COLOR': colors[i]}
            ball_urdf = self.fill_template(ball_urdf, replace)
            ball_id = env.add_object(ball_urdf, ball_pose)
            balls.append(ball_id)

        # Goal: each ball is in a different colored zone.
        for i in range(6):
            self.add_goal(objs=[balls[i]], matches=np.int32([[1]]), targ_poses=[zone_poses[i]], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 6)
            self.lang_goals.append(self.lang_template.format(color=color_names[i]))```
 END


Prompt:
 "color-coordinated-cylinder-tower". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "color-coordinated-cylinder-tower" in a python code block.
###

Response:
 color-coordinated-cylinder-tower: {'task-name': 'color-coordinated-cylinder-tower', 'task-description': 'Stack cylinders of four different colors (red, blue, green, yellow) on top of each other on a square stand in a specific sequence. The bottom of the stack should start with a blue cylinder, follow by a green cylinder, then a red one, and finally a yellow cylinder at the top. Each cylinder has to be aligned correctly to avoid falling.', 'assets-used': ['cylinder/cylinder-template.urdf', 'stacking/stand.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class ColorCoordinatedCylinderTower(Task):
    """Stack cylinders of four different colors (red, blue, green, yellow) on top of each other on a square stand in a specific sequence. The bottom of the stack should start with a blue cylinder, follow by a green cylinder, then a red one, and finally a yellow cylinder at the top. Each cylinder has to be aligned correctly to avoid falling."""

    def __init__(self):
        super().__init__()
        self.max_steps = 10
        self.lang_template = "Stack cylinders of four different colors (red, blue, green, yellow) on top of each other on a square stand in a specific sequence. The bottom of the stack should start with a blue cylinder, follow by a green cylinder, then a red one, and finally a yellow cylinder at the top."
        self.task_completed_desc = "done stacking cylinders."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add base.
        base_size = (0.05, 0.15, 0.005)
        base_urdf = 'stacking/stand.urdf'
        base_pose = self.get_random_pose(env, base_size)
        env.add_object(base_urdf, base_pose, category='fixed')

        # Cylinder colors.
        colors = [utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['red'], utils.COLORS['yellow']]

        # Add cylinders.
        cylinder_size = (0.04, 0.04, 0.04)
        cylinder_urdf = 'cylinder/cylinder-template.urdf'

        objs = []
        for i in range(4):
            cylinder_pose = self.get_random_pose(env, cylinder_size)
            cylinder_id = env.add_object(cylinder_urdf, cylinder_pose, color=colors[i])
            objs.append(cylinder_id)

        # Associate placement locations for goals.
        place_pos = [(0, 0, 0.03), (0, 0, 0.08), (0, 0, 0.13), (0, 0, 0.18)]
        targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]

        # Goal: cylinders are stacked in a tower (bottom to top: blue, green, red, yellow).
        for i in range(4):
            self.add_goal(objs=[objs[i]], matches=np.ones((1, 1)), targ_poses=[targs[i]], replace=False,
                    rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2])
            self.lang_goals.append(self.lang_template)```
 END


Prompt:
 "symmetric-block-bridge-construction". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "symmetric-block-bridge-construction" in a python code block.
###

Response:
 symmetric-block-bridge-construction: {'task-name': 'symmetric-block-bridge-construction', 'task-description': 'Create a symmetrical bridge-shaped structure on a stand using eight blocks of two different colors (four red and four blue). Start by placing two red blocks side by side at the center of the stand to form the base of the bridge. Then, take a blue block and place it on top of the red blocks, followed by another red block on top of the blue one, and this pattern continues till you exhaust all the blocks. The final structure should be a bridge with alternating colors (red, blue, red, blue). The challenge lies in ensuring symmetry and balancing the blocks without making them fall.', 'assets-used': ['stacking/stand.urdf', 'block/block-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
import pybullet as p

class SymmetricBlockBridgeConstruction(Task):
    """Create a symmetrical bridge-shaped structure on a stand using eight blocks of two different colors (four red and four blue)."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "create a symmetrical bridge-shaped structure on a stand using eight blocks of two different colors (four red and four blue)"
        self.task_completed_desc = "done building the bridge."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add base.
        base_size = (0.05, 0.15, 0.005)
        base_urdf = 'stacking/stand.urdf'
        base_pose = self.get_random_pose(env, base_size)
        env.add_object(base_urdf, base_pose, category='fixed')

        # Block colors.
        colors = [utils.COLORS['red'], utils.COLORS['blue']]

        # Add blocks.
        # x, y, z dimensions for the asset size
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'stacking/block.urdf'

        objs = []
        for i in range(8):
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=colors[i%2])
            objs.append(block_id)

        # Associate placement locations for goals.
        place_pos = [(0, -0.05, 0.03), (0, 0, 0.03),
                     (0, 0.05, 0.03), (0, -0.025, 0.08),
                     (0, 0.025, 0.08), (0, 0, 0.13),
                     (0, -0.025, 0.18), (0, 0.025, 0.18)]
        targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]

        # Goal: blocks are stacked in a bridge (bottom row: red, red).
        self.add_goal(objs=objs[:2], matches=np.ones((2, 2)), targ_poses=targs[:2], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2]*2)
        self.lang_goals.append(self.lang_template)

        # Goal: blocks are stacked in a bridge (second row: blue).
        self.add_goal(objs=objs[2:3], matches=np.ones((1, 1)), targ_poses=targs[2:3], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2])
        self.lang_goals.append(self.lang_template)

        # Goal: blocks are stacked in a bridge (third row: red).
        self.add_goal(objs=objs[3:4], matches=np.ones((1, 1)), targ_poses=targs[3:4], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2])
        self.lang_goals.append(self.lang_template)

        # Goal: blocks are stacked in a bridge (fourth row: blue).
        self.add_goal(objs=objs[4:5], matches=np.ones((1, 1)), targ_poses=targs[4:5], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2])
        self.lang_goals.append(self.lang_template)

        # Goal: blocks are stacked in a bridge (fifth row: red).
        self.add_goal(objs=objs[5:6], matches=np.ones((1, 1)), targ_poses=targs[5:6], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2])
        self.lang_goals.append(self.lang_template)

        # Goal: blocks are stacked in a bridge (sixth row: blue).
        self.add_goal(objs=objs[6:7], matches=np.ones((1, 1)), targ_poses=targs[6:7], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2])
        self.lang_goals.append(self.lang_template)

        # Goal: blocks are stacked in a bridge (top row: red, red).
        self.add_goal(objs=objs[7:], matches=np.ones((1, 1)), targ_poses=targs[7:], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2])
        self.lang_goals.append(self.lang_template)```
 END


Prompt:
 "sphere-align-stand". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "sphere-align-stand" in a python code block.
###

Response:
 sphere-align-stand: {'task-name': 'sphere-align-stand', 'task-description': 'On a table there are five differently colored stands and five spheres. The task involves picking up each sphere and placing it on the stand of the matching color. The task is challenging due to the precision required in picking up and placing the spheres, and the color coordination.', 'assets-used': ['stacking/stand.urdf', 'sphere/sphere.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class SphereAlignStand(Task):
    """Pick up each sphere and place it on the stand of the matching color."""

    def __init__(self):
        super().__init__()
        self.max_steps = 5
        self.lang_template = "place the {color} sphere on the {color} stand"
        self.task_completed_desc = "done aligning spheres with stands."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors for the spheres and stands
        colors = ['red', 'green', 'blue', 'yellow', 'purple']
        color_names = ['red', 'green', 'blue', 'yellow', 'purple']

        # Add stands.
        # x, y, z dimensions for the asset size
        stand_size = (0.05, 0.05, 0.05)
        stand_urdf = 'stacking/stand.urdf'
        stand_poses = []
        for i in range(5):
            stand_pose = self.get_random_pose(env, stand_size)
            env.add_object(stand_urdf, stand_pose, 'fixed', color=utils.COLORS[colors[i]])
            stand_poses.append(stand_pose)

        # Add spheres.
        # x, y, z dimensions for the asset size
        sphere_size = (0.04, 0.04, 0.04)
        sphere_urdf = 'sphere/sphere.urdf'
        spheres = []
        for i in range(5):
            sphere_pose = self.get_random_pose(env, sphere_size)
            sphere_id = env.add_object(sphere_urdf, sphere_pose, color=utils.COLORS[colors[i]])
            spheres.append(sphere_id)

        # Goal: each sphere is on the stand of the matching color.
        for i in range(5):
            self.add_goal(objs=[spheres[i]], matches=np.ones((1, 1)), targ_poses=[stand_poses[i]], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1/5)
            self.lang_goals.append(self.lang_template.format(color=color_names[i]))```
 END


Prompt:
 "construct-colorful-arch". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "construct-colorful-arch" in a python code block.
###

Response:
 construct-colorful-arch: {'task-name': 'construct-colorful-arch', 'task-description': 'Construct an arch using six blocks: three red, and three blue. Place two red blocks on the tabletop parallel to each other, then place a blue block on top of the red blocks to form a basic arch. Repeat the process by placing a red block on each side of the base arch and bridge them with the last blue block. Finally, place the last two red blocks on top of the second layer of the red blocks, on either side. The overall structure forms a colorful 3D arch.', 'assets-used': ['block/block.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class ConstructColorfulArch(Task):
    """Construct an arch using six blocks: three red, and three blue."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "Construct an arch using six blocks: three red, and three blue."
        self.task_completed_desc = "done constructing colorful arch."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add blocks.
        # x, y, z dimensions for the asset size
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'stacking/block.urdf'
        colors = [utils.COLORS['red'], utils.COLORS['blue']]
        blocks = []
        for i in range(6):
            block_pose = self.get_random_pose(env, block_size)
            color = colors[i // 3]  # First three blocks are red, last three are blue
            block_id = env.add_object(block_urdf, block_pose, color=color)
            blocks.append(block_id)

        # Associate placement locations for goals.
        place_pos = [(0, -0.05, 0.02), (0, 0.05, 0.02),  # Base layer
                     (0, 0, 0.06),  # Second layer
                     (0, -0.05, 0.10), (0, 0.05, 0.10),  # Third layer
                     (0, 0, 0.14)]  # Top layer
        targs = [(utils.apply(block_pose, i), block_pose[1]) for i in place_pos]

        # Goal: blocks are stacked in an arch (bottom layer: red, red).
        self.add_goal(objs=blocks[:2], matches=np.ones((2, 2)), targ_poses=targs[:2], replace=False,
                      rotations=True, metric='pose', params=None, step_max_reward=1 / 3, symmetries=[np.pi/2]*2)
        self.lang_goals.append("Place two red blocks on the tabletop parallel to each other")

        # Goal: blocks are stacked in an arch (second layer: blue).
        self.add_goal(objs=blocks[2:3], matches=np.ones((1, 1)), targ_poses=targs[2:3], replace=False,
                      rotations=True, metric='pose', params=None, step_max_reward=1 / 3, symmetries=[np.pi/2])
        self.lang_goals.append("Place a blue block on top of the red blocks to form a basic arch")

        # Goal: blocks are stacked in an arch (third layer: red, red).
        self.add_goal(objs=blocks[3:5], matches=np.ones((2, 2)), targ_poses=targs[3:5], replace=False,
                      rotations=True, metric='pose', params=None, step_max_reward=1 / 3, symmetries=[np.pi/2]*2)
        self.lang_goals.append("Place a red block on each side of the base arch")

        # Goal: blocks are stacked in an arch (top layer: blue).
        self.add_goal(objs=blocks[5:], matches=np.ones((1, 1)), targ_poses=targs[5:], replace=False,
                      rotations=True, metric='pose', params=None, step_max_reward=1 / 3, symmetries=[np.pi/2])
        self.lang_goals.append("Bridge them with the last blue block")```
 END


Prompt:
 "color-sorted-container-stack". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "color-sorted-container-stack" in a python code block.
###

Response:
 color-sorted-container-stack: {'task-name': 'color-sorted-container-stack', 'task-description': 'The challenge is to stack four differently colored blocks (red, blue, green, yellow) inside a container. These blocks are initially placed inside two separate bowls. The red and blue blocks are in one bowl, the yellow and green blocks are in another. The robot has to pick up the blocks from the two bowls and stack them inside the container in the specific order: red, blue, green, then yellow.', 'assets-used': ['bowl/bowl.urdf', 'container/container-template.urdf', 'block/block.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class ColorSortedContainerStack(Task):
    """Stack four differently colored blocks (red, blue, green, yellow) inside a container."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "stack the blocks in the container in the order: red, blue, green, then yellow"
        self.task_completed_desc = "done stacking blocks."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add container.
        # x, y, z dimensions for the asset size
        container_size = (0.15, 0.15, 0.15)
        container_pose = self.get_random_pose(env, container_size)
        container_urdf = 'container/container-template.urdf'
        replace = {'DIM': container_size, 'HALF': (container_size[0] / 2, container_size[1] / 2, container_size[2] / 2)}
        container_urdf = self.fill_template(container_urdf, replace)
        env.add_object(container_urdf, container_pose, 'fixed')

        # Add blocks.
        # x, y, z dimensions for the asset size
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        block_colors = [utils.COLORS['red'], utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['yellow']]
        blocks = []
        for i in range(4):
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=block_colors[i])
            blocks.append(block_id)

        # Add bowls.
        # x, y, z dimensions for the asset size
        bowl_size = (0.12, 0.12, 0)
        bowl_urdf = 'bowl/bowl.urdf'
        for i in range(2):
            bowl_pose = self.get_random_pose(env, bowl_size)
            env.add_object(bowl_urdf, bowl_pose, 'fixed')

        # Goal: each block is stacked in the container in the order: red, blue, green, yellow.
        for i in range(4):
            self.add_goal(objs=[blocks[i]], matches=np.ones((1, 1)), targ_poses=[container_pose], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 4)
        self.lang_goals.append(self.lang_template)```
 END


Prompt:
 "align-spheres-in-colored-zones". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "align-spheres-in-colored-zones" in a python code block.
###

Response:
 align-spheres-in-colored-zones: {'task-name': 'align-spheres-in-colored-zones', 'task-description': 'There are four spheres of different colors (red, blue, green, yellow) positioned randomly on the table along with four zones marked with matching colors. The task is to pick up each sphere and place it into the matching colored zone with precise placement.', 'assets-used': ['sphere/sphere-template.urdf', 'zone/zone.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
import pybullet as p

class AlignSpheresInColoredZones(Task):
    """Align spheres of different colors in the matching colored zones."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "place the {color} sphere in the {color} zone"
        self.task_completed_desc = "done aligning spheres in colored zones."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors
        colors = ['red', 'blue', 'green', 'yellow']
        color_names = ['red', 'blue', 'green', 'yellow']

        # Add zones.
        zone_size = (0.12, 0.12, 0)
        zone_urdf = 'zone/zone.urdf'
        zone_poses = []
        for color in colors:
            zone_pose = self.get_random_pose(env, zone_size)
            env.add_object(zone_urdf, zone_pose, 'fixed', color=utils.COLORS[color])
            zone_poses.append(zone_pose)

        # Add spheres.
        sphere_size = (0.04, 0.04, 0.04)
        sphere_urdf = 'sphere/sphere-template.urdf'
        spheres = []
        for i, color in enumerate(colors):
            sphere_pose = self.get_random_pose(env, sphere_size)
            replace = {'DIM': sphere_size, 'HALF': (sphere_size[0] / 2, sphere_size[1] / 2, sphere_size[2] / 2)}
            sphere_urdf = self.fill_template(sphere_urdf, replace)
            sphere_id = env.add_object(sphere_urdf, sphere_pose, color=utils.COLORS[color])
            spheres.append(sphere_id)

            # Add goal
            self.add_goal(objs=[sphere_id], matches=np.ones((1, 1)), targ_poses=[zone_poses[i]], replace=False,
                          rotations=False, metric='pose', params=None, step_max_reward=1)
            self.lang_goals.append(self.lang_template.format(color=color_names[i]))```
 END


Prompt:
 "sort-insert-color-coordinated-blocks". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "sort-insert-color-coordinated-blocks" in a python code block.
###

Response:
 sort-insert-color-coordinated-blocks: {'task-name': 'sort-insert-color-coordinated-blocks', 'task-description': 'There are six blocks of three different colors (two red, two blue, and two green), two red containers, two blue containers, and two green containers scattered on the table. The task is to sort the blocks by their colors and place (insert) them into the containers of the matching color. Two blocks of the same color should go into two different containers of the same color. The challenge lies in the color-based sorting and insertion of the blocks into the containers.', 'assets-used': ['block/block.urdf', 'container/container-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class SortInsertColorCoordinatedBlocks(Task):
    """Sort blocks by their colors and place them into the containers of the matching color."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "sort the blocks by their colors and place them into the containers of the matching color"
        self.task_completed_desc = "done sorting and inserting blocks."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add containers.
        container_size = (0.12, 0.12, 0.12)
        container_urdf = 'container/container-template.urdf'
        container_colors = ['red', 'blue', 'green']
        container_poses = []
        for color in container_colors:
            container_pose = self.get_random_pose(env, container_size)
            env.add_object(container_urdf, container_pose, color=utils.COLORS[color])
            container_poses.append(container_pose)

        # Add blocks.
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        block_colors = ['red', 'red', 'blue', 'blue', 'green', 'green']
        blocks = []
        for color in block_colors:
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=utils.COLORS[color])
            blocks.append(block_id)

        # Goal: each block is in a container of the same color.
        for i in range(len(blocks)):
            self.add_goal(objs=[blocks[i]], matches=np.ones((1, 1)), targ_poses=[container_poses[i//2]], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1/len(blocks))
        self.lang_goals.append(self.lang_template)```
 END


Prompt:
 "color-ordered-insertion". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "color-ordered-insertion" in a python code block.
###

Response:
 color-ordered-insertion: {'task-name': 'color-ordered-insertion', 'task-description': 'There are four differently-colored ell objects (red, blue, green, yellow) and a corresponding set of color-coded fixtures. The task involves picking up each ell object and inserting it into the matching color fixture in a specific order: from left to right, insert red, blue, green, and finally yellow. The challenge lies in the precise manipulation of the ell objects and the color-coordination required. The fixtures are arranged in a straight line, and can only be approached from one direction, demanding careful navigation.', 'assets-used': ['insertion/ell.urdf', 'insertion/fixture.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class ColorOrderedInsertion(Task):
    """Insert differently-colored ell objects into the matching color fixture in a specific order."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "insert the {color} ell into the matching fixture"
        self.task_completed_desc = "done inserting ells."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors and their order
        colors = ['red', 'blue', 'green', 'yellow']
        color_order = {color: i for i, color in enumerate(colors)}

        # Add fixtures.
        fixture_size = (0.12, 0.12, 0.02)
        fixture_urdf = 'insertion/fixture.urdf'
        fixtures = []
        for color in colors:
            fixture_pose = self.get_random_pose(env, fixture_size)
            fixture_id = env.add_object(fixture_urdf, fixture_pose, color=utils.COLORS[color], category='fixed')
            fixtures.append(fixture_id)

        # Add ell objects.
        ell_size = (0.04, 0.04, 0.04)
        ell_urdf = 'insertion/ell.urdf'
        ells = []
        for color in colors:
            ell_pose = self.get_random_pose(env, ell_size)
            ell_id = env.add_object(ell_urdf, ell_pose, color=utils.COLORS[color])
            ells.append(ell_id)

        # Goal: each ell is inserted into the matching color fixture in the correct order.
        for i, ell in enumerate(ells):
            self.add_goal(objs=[ell], matches=np.ones((1, 1)), targ_poses=[p.getBasePositionAndOrientation(fixtures[i])], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1 / len(ells))
            self.lang_goals.append(self.lang_template.format(color=colors[i]))```
 END


Prompt:
 "color-coordinated-cylinders-in-boxes". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "color-coordinated-cylinders-in-boxes" in a python code block.
###

Response:
 color-coordinated-cylinders-in-boxes: {'task-name': 'color-coordinated-cylinders-in-boxes', 'task-description': 'Arrange four cylinders of different colors (red, blue, green, and yellow) into four matching colored boxes, such that each box contains a cylinder of the same color. The task is challenging due to the manipulation of cylindrical objects and the color coordination required.', 'assets-used': ['cylinder/cylinder-template.urdf', 'box/box-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class ColorCoordinatedCylindersInBoxes(Task):
    """Arrange four cylinders of different colors (red, blue, green, and yellow) into four matching colored boxes, such that each box contains a cylinder of the same color."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "put the {color} cylinder in the {color} box"
        self.task_completed_desc = "done arranging cylinders in boxes."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors and corresponding names
        colors = [utils.COLORS['red'], utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['yellow']]
        color_names = ['red', 'blue', 'green', 'yellow']

        # Add boxes.
        box_size = (0.05, 0.05, 0.05)  # x, y, z dimensions for the box size
        box_urdf = 'box/box-template.urdf'
        boxes = []
        for i in range(4):
            box_pose = self.get_random_pose(env, box_size)
            box_id = env.add_object(box_urdf, box_pose, color=colors[i])
            boxes.append(box_id)

        # Add cylinders.
        cylinder_size = (0.02, 0.02, 0.05)  # x, y, z dimensions for the cylinder size
        cylinder_urdf = 'cylinder/cylinder-template.urdf'
        cylinders = []
        for i in range(4):
            cylinder_pose = self.get_random_pose(env, cylinder_size)
            cylinder_id = env.add_object(cylinder_urdf, cylinder_pose, color=colors[i])
            cylinders.append(cylinder_id)

        # Goal: each cylinder is in a box of the same color.
        for i in range(4):
            self.add_goal(objs=[cylinders[i]], matches=np.ones((1, 1)), targ_poses=[box_pose], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1 / 4)
            self.lang_goals.append(self.lang_template.format(color=color_names[i]))```
 END


Prompt:
 "color-coordinated-insertion". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "color-coordinated-insertion" in a python code block.
###

Response:
 color-coordinated-insertion: {'task-name': 'color-coordinated-insertion', 'task-description': 'There are three insertion fixtures and three ell shaped blocks of different colors (red, blue, green) on the table top. The task is to pick up the ell shaped blocks and insert each one of them into the fixture of the same color. However, the ell blocks should be inserted in a specific sequence - red first, then blue, and finally green. This task is challenging due to the precision required for insertion and the need for color coordination and sequencing.', 'assets-used': ['insertion/ell.urdf', 'insertion/fixture.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
import pybullet as p

class ColorCoordinatedInsertion(Task):
    """Insert each block into the fixture of the same color, and then place these fixtures onto a pallet."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "insert each block into the fixture of the same color, and then place these fixtures onto a pallet"
        self.task_completed_desc = "done with color-coordinated-insertion."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add pallet.
        pallet_size = (0.35, 0.35, 0.01)
        pallet_pose = self.get_random_pose(env, pallet_size)
        pallet_urdf = 'pallet/pallet.urdf'
        env.add_object(pallet_urdf, pallet_pose, 'fixed')

        # Add fixtures and blocks.
        colors = ['red', 'blue', 'green', 'yellow']
        fixtures = []
        blocks = []
        fixture_size = (0.05, 0.05, 0.05)
        block_size = (0.04, 0.04, 0.04)
        fixture_urdf = 'insertion/fixture.urdf'
        block_urdf = 'block/block.urdf'
        for color in colors:
            # Add fixture.
            fixture_pose = self.get_random_pose(env, fixture_size)
            fixture_id = env.add_object(fixture_urdf, fixture_pose, color=utils.COLORS[color])
            fixtures.append(fixture_id)

            # Add block.
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=utils.COLORS[color])
            blocks.append(block_id)

        # Goal: each block is in the fixture of the same color.
        for i in range(len(blocks)):
            self.add_goal(objs=[blocks[i]], matches=np.ones((1, 1)), targ_poses=[p.getBasePositionAndOrientation(fixtures[i])], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1 / len(blocks))

        # Goal: each fixture is on the pallet.
        for i in range(len(fixtures)):
            self.add_goal(objs=[fixtures[i]], matches=np.ones((1, 1)), targ_poses=[pallet_pose], replace=False,
                          rotations=True, metric='zone', params=[(pallet_pose, pallet_size)], step_max_reward=1 / len(fixtures))

        self.lang_goals.append(self.lang_template)```
 END


Prompt:
 "cylinder-stand-alignment". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "cylinder-stand-alignment" in a python code block.
###

Response:
 cylinder-stand-alignment: {'task-name': 'cylinder-stand-alignment', 'task-description': 'Arrange four colored cylinders (red, blue, green, yellow) in order of their colors on four stands of matching color. However, the stands are placed in a random order on the table, which increases the complexity of the task as it requires careful planning and color matching.', 'assets-used': ['cylinder/cylinder-template.urdf', 'stacking/stand.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class CylinderStandAlignment(Task):
    """Arrange four colored cylinders (red, blue, green, yellow) in order of their colors on four stands of matching color."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "Arrange the {color} cylinder on the {color} stand"
        self.task_completed_desc = "done arranging cylinders."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors and corresponding names
        colors = [utils.COLORS['red'], utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['yellow']]
        color_names = ['red', 'blue', 'green', 'yellow']

        # Add cylinders.
        # x, y, z dimensions for the asset size
        cylinder_size = (0.04, 0.04, 0.04)
        cylinder_urdf = 'cylinder/cylinder-template.urdf'
        cylinders = []
        for i in range(4):
            cylinder_pose = self.get_random_pose(env, cylinder_size)
            replace = {'DIM': cylinder_size, 'HALF': (cylinder_size[0] / 2, cylinder_size[1] / 2, cylinder_size[2] / 2),
                       'COLOR': colors[i]}
            # IMPORTANT: REPLACE THE TEMPLATE URDF
            urdf = self.fill_template(cylinder_urdf, replace)
            cylinder_id = env.add_object(urdf, cylinder_pose)
            cylinders.append(cylinder_id)

        # Add stands.
        # x, y, z dimensions for the asset size
        stand_size = (0.05, 0.05, 0.005)
        stand_urdf = 'stacking/stand.urdf'
        stands = []
        for i in range(4):
            stand_pose = self.get_random_pose(env, stand_size)
            env.add_object(stand_urdf, stand_pose, color=colors[i], category='fixed')
            stands.append(stand_pose)

        # Goal: each cylinder is on a stand of the same color.
        for i in range(4):
            self.add_goal(objs=[cylinders[i]], matches=np.ones((1, 1)), targ_poses=[stands[i]], replace=False,
                          rotations=True, metric='pose', params=None, step_max_reward=1 / 4)
            self.lang_goals.append(self.lang_template.format(color=color_names[i]))```
 END


Prompt:
 "color-coordinated-container-sorting". Format the answer in a python dictionary with keys "task-name", "task-description", and "assets-used". Then write the pybullet simulation code for the task "color-coordinated-container-sorting" in a python code block.
###

Response:
 color-coordinated-container-sorting: {'task-name': 'color-coordinated-container-sorting', 'task-description': 'On the table there are four different colored containers (red, blue, green and yellow) and sixteen blocks of the same four colors (four blocks of each color). The task challenges the robot to pick up these blocks and sort them into the containers of the same color, with each container having one of each color of blocks in it. To make the task interesting, each container should have the blocks in a different order. One container might have them in the order red, blue, green, yellow; another one in the order yellow, green, blue, red; and the last two could have them in the orders green, red, yellow, blue; and blue, yellow, red, green. The robot would need to first identify the correct container for sorting, and then plan the correct order of the blocks.', 'assets-used': ['block/block.urdf', 'container/container-template.urdf']}```

###Implementation:
  ```python
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class ColorCoordinatedContainerSorting(Task):
    """Sort blocks into containers of the same color in a specific order."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "sort the blocks into the containers of the same color in the specified order"
        self.task_completed_desc = "done sorting blocks into containers."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Define colors and their orders in containers
        colors = ['red', 'blue', 'green', 'yellow']
        orders = [
            ['red', 'blue', 'green', 'yellow'],
            ['yellow', 'green', 'blue', 'red'],
            ['green', 'red', 'yellow', 'blue'],
            ['blue', 'yellow', 'red', 'green']
        ]

        # Add containers.
        container_size = (0.12, 0.12, 0.02)
        container_urdf = 'container/container-template.urdf'
        containers = []
        for i in range(4):
            container_pose = self.get_random_pose(env, container_size)
            color = utils.COLORS[colors[i]]
            container_id = env.add_object(container_urdf, container_pose, color=color)
            containers.append(container_id)

        # Add blocks.
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'
        blocks = []
        for i in range(4):
            for _ in range(4):
                block_pose = self.get_random_pose(env, block_size)
                color = utils.COLORS[colors[i]]
                block_id = env.add_object(block_urdf, block_pose, color=color)
                blocks.append(block_id)

        # Add goals.
        for i in range(4):
            for j in range(4):
                # Find the block and container of the same color
                block_id = blocks[i*4 + j]
                container_id = containers[i]
                # Define the target pose based on the order in the container
                container_pose = p.getBasePositionAndOrientation(container_id)[0]
                targ_pose = (container_pose[0], container_pose[1], container_pose[2] + block_size[2] * (j + 1))
                self.add_goal(objs=[block_id], matches=np.ones((1, 1)), targ_poses=[targ_pose], replace=False,
                              rotations=True, metric='pose', params=None, step_max_reward=1/16)

        self.lang_goals.append(self.lang_template)```
 END