Expert Code Generation¶

1. Code_gen Folder Structure¶

This directory contains various modules for generating and testing robot task code:

• gpt_agent.py: API integration with LLM models
• observation_agent.py: Processes multi-modal observations for code correction
• prompt.py: Prompt templates for code generation
• run_code.py: Executes and tests generated code
• task_generation_simple.py: Basic single-pass code generation
• task_generation.py: Iterative code generation with error feedback
• task_generation_mm.py: Advanced code generation with multi-modal observation
• task_info.py: Task definitions and descriptions
• test_gen_code.py: Utility for testing generated code with detailed metrics

The code generation system also interacts with these important directories: - ./envs/: Contains manually implemented task environments - _base_task.py: Core environment with robot control functions and utilities - Includes save_camera_images(task_name, step_name, generate_num_id, save_dir) for capturing visual observations during task execution - ./envs_gen/: Stores auto-generated task implementations - ./task_config/: Configuration files for tasks and embodiments - ./script/: Template scripts and utilities - ./assets/objects/: 3D models and metadata for simulation objects - ./camera_images/: Stores observation images captured during code generation for multi-modal feedback

The entire pipeline enables automatic generation of robot control code from natural language task descriptions, with feedback-based refinement and multi-modal observation capabilities.

2. Configure LLM API Key¶

Please configure the necessary API keys in the code_gen/gpt_agent.py file. Additionally, if the LLM you are utilizing does not support integration with the OpenAI API, you may need to make corresponding adjustments to the generate() function.

3. Generate Your Task Code¶

3.1 1. Add Task Description¶

Add new task information, including the task name and natural language description, in ./code_gen/task_info.py.

1. Template of Task Information:¶

TASK_NAME = {
    "task_name": "task_name",                # Name of the task
    "task_description": "...",               # Detailed description of the task
    "current_code": '''
                class gpt_{task_name}({task_name}):
                    def play_once(self):
                        pass
                '''                          # Code template to be completed
    "actor_list": {                          # List of involved objects; can be a dictionary or a simple list
        "self.object1": {
            "name": "object1",               # Object name
            "description": "...",            # Description of the object
            "modelname": "model_name"        # Name of the 3D model representing the object
        },
        "self.object2": {
            "name": "object2",
            "description": "...",
            "modelname": "model_name"
        },
        # ... more objects
    },
    # Alternatively, the actor_list can be a simple list:
    # "actor_list": ["self.object1", "self.object2", ...],
    # To make code generation easier, the actor_list also includes some pose information
    # like target poses or middle poses (optional and don't require modelname).
}

1.1 2. Add Basic Task Code¶

Add the basic code file ${task_name}.py in the ./envs/ directory, following this structure:

from .base_task import Base_task
from .utils import *
import sapien

class ${task_name}(Base_task):
    def setup_demo(self, **kwargs):
        # Initializes the simulation environment for the task
        # Sets up the table, robot, planner, camera, and initial positions
        # This function is called once at the beginning of each episode
        pass

    def load_actors(self):
        # Loads all the necessary objects for the task into the environment
        # Typically called from setup_demo to initialize scene objects
        # Can also be used to set initial poses for objects
        pass

    def play_once(self):
        # Contains the robot control code to complete the task
        # This is the main function that will be generated by the LLM
        # Implements the sequence of actions for the robot to achieve the task
        pass

    # Check success
    def check_success(self):
        # Defines criteria to determine if the task was completed successfully
        # Returns a boolean indicating success or failure
        # Used for evaluation and feedback during code generation
        pass

In the code above, {task_name} should match the name of the basic code file, and the check_success() function is used to determine if the task is successful. No changes are needed for the rest of the code.

Note: The envs folder contains manually written files with setup_demo, robot operation code in play_once, and check_success methods. Auto-generated code will be saved in the envs_gen folder.

1.2 3. Generate the Final Code¶

You can use three different code generation approaches depending on your needs:

Note: The code generation process will only generate the play_once() method implementation, which contains the robot control logic to complete the task. Other methods like setup_demo(), load_actors(), and check_success() should be manually implemented.

1.2.1 Basic Code Generation¶

For quick verification of new tasks or debugging existing ones without iterative correction:

python code_gen/task_generation_simple.py task_name

1.2.2 Code Generation with Error Feedback¶

This script implements iterative code correction based on error feedback, consistent with RoboTwin 1.0:

python code_gen/task_generation.py task_name

1.2.3 Advanced Code Generation with Multi-Modal Observations¶

This script provides both error feedback iteration and multi-modal observation-based code correction, consistent with RoboTwin 2.0. It offers the best generation quality but runs slower:

python code_gen/task_generation_mm.py task_name

The multi-modal observation functionality is implemented in code_gen/observation_agent.py.

The generated code file will be saved as ./envs_gen/gpt_${task_name}.py. For example:

python code_gen/task_generation_mm.py pick_dual_bottles_easy

1.3 4. Test Generated Code¶

Run the following script to test the generated code:

python code_gen/run_code.py task_name

This will execute the task using the generated code and display the results, allowing you to validate the performance.

1.1 Additional Resources¶

For more information on generating task descriptions and object descriptions, refer to the documentation in the description directory.

For policy training and evaluation using the generated code, consult the policy/ACT documentation.