Adicionar suporte para multiplas estruturas (duas piramides): codigo exemplo e instrucoes nos prompts

cliport/generated_tasks/build_two_pyramids_blue_red_fixed.py

@@ -0,0 +1,158 @@
+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
+
+
+class BuildTwoPyramidsBlueRed(Task):
+    """Build two pyramids: one with blue blocks and one with red blocks."""
+
+    def __init__(self):
+        super().__init__()
+        self.max_steps = 20
+        self.lang_template = "build two pyramids: one with blue blocks and one with red blocks"
+        self.task_completed_desc = "done building two pyramids."
+        self.additional_reset()
+
+    def reset(self, env):
+        super().reset(env)
+
+        # Add two bases/stands for the pyramids
+        base_size = (0.05, 0.15, 0.005)
+        base_urdf = 'stacking/stand.urdf'
+        
+        # First base (for blue pyramid) - positioned on the left side of workspace
+        # Workspace bounds: X: [0.25, 0.75], Y: [-0.5, 0.5], Z: [0, 0.3]
+        base1_center = (0.4, -0.2, 0.001)  # Left side, within bounds
+        base1_rotation = (0, 0, 0, 1)
+        base1_pose = (base1_center, base1_rotation)
+        env.add_object(base_urdf, base1_pose, category='fixed')
+        
+        # Second base (for red pyramid) - positioned on the right side of workspace
+        base2_center = (0.6, 0.2, 0.001)  # Right side, within bounds
+        base2_rotation = (0, 0, 0, 1)
+        base2_pose = (base2_center, base2_rotation)
+        env.add_object(base_urdf, base2_pose, category='fixed')
+
+        # Add blocks: 6 blue blocks for first pyramid, 6 red blocks for second pyramid
+        block_size = (0.04, 0.04, 0.04)
+        block_urdf = 'stacking/block.urdf'
+        
+        blue_blocks = []
+        red_blocks = []
+        
+        # Add 6 blue blocks
+        for _ in range(6):
+            block_pose = self.get_random_pose(env, block_size)
+            block_id = env.add_object(block_urdf, block_pose, color=utils.COLORS['blue'])
+            blue_blocks.append(block_id)
+        
+        # Add 6 red blocks
+        for _ in range(6):
+            block_pose = self.get_random_pose(env, block_size)
+            block_id = env.add_object(block_urdf, block_pose, color=utils.COLORS['red'])
+            red_blocks.append(block_id)
+
+        # Define target poses for pyramids using RELATIVE positions
+        # Pyramid structure: 3 blocks base, 2 blocks middle, 1 block top
+        place_pos = [
+            (0, -0.05, 0.03),  # Left base
+            (0, 0.0, 0.03),    # Center base
+            (0, 0.05, 0.03),   # Right base
+            (0, -0.025, 0.08), # Left middle
+            (0, 0.025, 0.08),  # Right middle
+            (0, 0.0, 0.13)     # Top
+        ]
+        
+        # Apply relative positions to each base pose
+        blue_targs = [(utils.apply(base1_pose, pos), base1_pose[1]) for pos in place_pos]
+        red_targs = [(utils.apply(base2_pose, pos), base2_pose[1]) for pos in place_pos]
+
+        # Goal: Build blue pyramid (bottom row: 3 blocks)
+        self.add_goal(
+            objs=blue_blocks[:3], 
+            matches=np.ones((3, 3)), 
+            targ_poses=blue_targs[:3], 
+            replace=False,
+            rotations=True, 
+            metric='pose', 
+            params=None, 
+            step_max_reward=1.0 / 6, 
+            symmetries=[np.pi/2]*3,
+            language_goal="build the bottom row of the blue pyramid with three blue blocks"
+        )
+
+        # Goal: Build blue pyramid (middle row: 2 blocks)
+        self.add_goal(
+            objs=blue_blocks[3:5], 
+            matches=np.ones((2, 2)), 
+            targ_poses=blue_targs[3:5], 
+            replace=False,
+            rotations=True, 
+            metric='pose', 
+            params=None, 
+            step_max_reward=1.0 / 6, 
+            symmetries=[np.pi/2]*2,
+            language_goal="build the middle row of the blue pyramid with two blue blocks"
+        )
+
+        # Goal: Build blue pyramid (top: 1 block)
+        self.add_goal(
+            objs=blue_blocks[5:], 
+            matches=np.ones((1, 1)), 
+            targ_poses=blue_targs[5:], 
+            replace=False,
+            rotations=True, 
+            metric='pose', 
+            params=None, 
+            step_max_reward=1.0 / 6, 
+            symmetries=[np.pi/2],
+            language_goal="place the top block to complete the blue pyramid"
+        )
+
+        # Goal: Build red pyramid (bottom row: 3 blocks)
+        self.add_goal(
+            objs=red_blocks[:3], 
+            matches=np.ones((3, 3)), 
+            targ_poses=red_targs[:3], 
+            replace=False,
+            rotations=True, 
+            metric='pose', 
+            params=None, 
+            step_max_reward=1.0 / 6, 
+            symmetries=[np.pi/2]*3,
+            language_goal="build the bottom row of the red pyramid with three red blocks"
+        )
+
+        # Goal: Build red pyramid (middle row: 2 blocks)
+        self.add_goal(
+            objs=red_blocks[3:5], 
+            matches=np.ones((2, 2)), 
+            targ_poses=red_targs[3:5], 
+            replace=False,
+            rotations=True, 
+            metric='pose', 
+            params=None, 
+            step_max_reward=1.0 / 6, 
+            symmetries=[np.pi/2]*2,
+            language_goal="build the middle row of the red pyramid with two red blocks"
+        )
+
+        # Goal: Build red pyramid (top: 1 block)
+        self.add_goal(
+            objs=red_blocks[5:], 
+            matches=np.ones((1, 1)), 
+            targ_poses=red_targs[5:], 
+            replace=False,
+            rotations=True, 
+            metric='pose', 
+            params=None, 
+            step_max_reward=1.0 / 6, 
+            symmetries=[np.pi/2],
+            language_goal="place the top block to complete the red pyramid"
+        )
+

prompts/topdown_task_generation_prompt/cliport_prompt_code_split_template.txt

@@ -279,4 +279,12 @@ CRITICAL FOR PYRAMIDS, BRIDGES, AND MULTI-LAYER STRUCTURES:
 - Split the reward across layers: `step_max_reward=1.0 / num_layers` for each layer
 - Each layer should have its own `language_goal` describing that specific step (e.g., "build the bottom row", "build the middle row", "place the top block")
 
+CRITICAL FOR MULTIPLE STRUCTURES (e.g., "two pyramids", "three towers", "multiple zones"):
+- If the task requires MULTIPLE separate structures (e.g., "two pyramids", "one blue pyramid and one red pyramid"), you MUST create separate base poses for each structure
+- Use different base poses (e.g., `base1_pose`, `base2_pose`) positioned at different locations in the workspace
+- Create separate target pose lists for each structure using `utils.apply()` with each respective base pose
+- Use MULTIPLE sets of `add_goal()` calls - one set for each structure (e.g., 3 goals for first pyramid, 3 goals for second pyramid)
+- Each structure should have its own objects and its own goals, clearly separated in the code
+- Example: For "two pyramids", create `blue_blocks` and `red_blocks`, `blue_targs` and `red_targs`, and separate `add_goal()` calls for each pyramid
+
 Now write the code for the task "TASK_NAME_TEMPLATE" in python code block starting with ```python.  Reminder: TASK_STRING_TEMPLATE 

prompts/topdown_task_generation_prompt_simple/cliport_prompt_code_split_template.txt

@@ -277,4 +277,12 @@ CRITICAL FOR PYRAMIDS, BRIDGES, AND MULTI-LAYER STRUCTURES:
 - Split the reward across layers: `step_max_reward=1.0 / num_layers` for each layer
 - Each layer should have its own `language_goal` describing that specific step (e.g., "build the bottom row", "build the middle row", "place the top block")
 
+CRITICAL FOR MULTIPLE STRUCTURES (e.g., "two pyramids", "three towers", "multiple zones"):
+- If the task requires MULTIPLE separate structures (e.g., "two pyramids", "one blue pyramid and one red pyramid"), you MUST create separate base poses for each structure
+- Use different base poses (e.g., `base1_pose`, `base2_pose`) positioned at different locations in the workspace
+- Create separate target pose lists for each structure using `utils.apply()` with each respective base pose
+- Use MULTIPLE sets of `add_goal()` calls - one set for each structure (e.g., 3 goals for first pyramid, 3 goals for second pyramid)
+- Each structure should have its own objects and its own goals, clearly separated in the code
+- Example: For "two pyramids", create `blue_blocks` and `red_blocks`, `blue_targs` and `red_targs`, and separate `add_goal()` calls for each pyramid
+
 Now write the code for the task "TASK_NAME_TEMPLATE" in python code block starting with ```python.  Reminder: TASK_STRING_TEMPLATE 

prompts/vanilla_task_generation_prompt/cliport_prompt_code_split_template.txt

@@ -266,6 +266,14 @@ CRITICAL FOR PYRAMIDS, BRIDGES, AND MULTI-LAYER STRUCTURES:
 - Include `symmetries=[np.pi/2]*n` parameter in `add_goal()` to allow 90-degree rotations
 - Split the reward across layers: `step_max_reward=1.0 / num_layers` for each layer
 - Each layer should have its own `language_goal` describing that specific step (e.g., "build the bottom row", "build the middle row", "place the top block")
+
+CRITICAL FOR MULTIPLE STRUCTURES (e.g., "two pyramids", "three towers", "multiple zones"):
+- If the task requires MULTIPLE separate structures (e.g., "two pyramids", "one blue pyramid and one red pyramid"), you MUST create separate base poses for each structure
+- Use different base poses (e.g., `base1_pose`, `base2_pose`) positioned at different locations in the workspace
+- Create separate target pose lists for each structure using `utils.apply()` with each respective base pose
+- Use MULTIPLE sets of `add_goal()` calls - one set for each structure (e.g., 3 goals for first pyramid, 3 goals for second pyramid)
+- Each structure should have its own objects and its own goals, clearly separated in the code
+- Example: For "two pyramids", create `blue_blocks` and `red_blocks`, `blue_targs` and `red_targs`, and separate `add_goal()` calls for each pyramid
 - ALL target poses MUST be within workspace bounds: X: [0.25, 0.75], Y: [-0.5, 0.5], Z: [0, 0.3]
 
 Now write the code for the task "TASK_NAME_TEMPLATE" in python code block starting with ```python.  Reminder: TASK_STRING_TEMPLATE