File size: 5,034 Bytes
13b4740
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
"""
Sistema de templates flexíveis para estruturas comuns.
Templates definem estrutura básica, LLM preenche lógica específica.
"""
from typing import Dict, Optional, List
import re


class CodeTemplate:
    """Template para estruturas comuns de código"""
    
    TEMPLATES = {
        'row': {
            'description': 'Template para criar fileira de blocos',
            'structure': """
# Add blocks.
block_size = (0.04, 0.04, 0.04)
block_urdf = 'stacking/block.urdf'
block_color = utils.COLORS['{color}']

objs = []
for _ in range({count}):
    block_pose = self.get_random_pose(env, block_size)
    block_id = env.add_object(block_urdf, block_pose, color=block_color)
    objs.append(block_id)

# Define target positions for the blocks in a row.
base_pose = ((0.5, 0.0, 0.02), (0, 0, 0, 1))
offsets = {offsets_placeholder}  # LLM generates: e.g., [(-0.05, 0, 0.02), (0, 0, 0.02), (0.05, 0, 0.02)]
targs = [(utils.apply(base_pose, offset), base_pose[1]) for offset in offsets]

# Goal: blocks are arranged in a row.
self.add_goal(objs=objs, matches=np.eye({count}), targ_poses=targs, replace=False,
              rotations=True, metric='pose', params=None, step_max_reward=1.0, language_goal=self.lang_template)
"""
        },
        'stack': {
            'description': 'Template para empilhar blocos',
            'structure': """
# Add blocks.
block_size = (0.04, 0.04, 0.04)
block_urdf = 'stacking/block.urdf'
block_color = utils.COLORS['{color}']

objs = []
for _ in range({count}):
    block_pose = self.get_random_pose(env, block_size)
    block_id = env.add_object(block_urdf, block_pose, color=block_color)
    objs.append(block_id)

# Define target positions for stacking.
base_pose = ((0.5, 0.0, 0.02), (0, 0, 0, 1))
stack_poses = {stack_poses_placeholder}  # LLM generates: e.g., [(0, 0, 0.04), (0, 0, 0.08), (0, 0, 0.12)]
targs = [(utils.apply(base_pose, pos), base_pose[1]) for pos in stack_poses]

# Goal: blocks are stacked.
self.add_goal(objs=objs, matches=np.eye({count}), targ_poses=targs, replace=False,
              rotations=True, metric='pose', params=None, step_max_reward=1.0, language_goal=self.lang_template)
"""
        },
        'pyramid': {
            'description': 'Template para pirâmide de blocos',
            'structure': """
# 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')

# Add blocks.
block_size = (0.04, 0.04, 0.04)
block_urdf = 'stacking/block.urdf'
block_color = utils.COLORS['{color}']

objs = []
for _ in range({count}):
    block_pose = self.get_random_pose(env, block_size)
    block_id = env.add_object(block_urdf, block_pose, color=block_color)
    objs.append(block_id)

# IMPORTANT: Associate placement locations for goals.
# Pyramid structure: {pyramid_structure_placeholder}  # LLM describes structure
place_pos = {place_pos_placeholder}  # LLM generates: e.g., [(0, -0.05, 0.03), (0, 0, 0.03), (0, 0.05, 0.03), ...]
targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]

# Goal: blocks are stacked in a pyramid.
# Use multiple add_goal() calls for each layer
{goals_placeholder}  # LLM generates multiple add_goal() calls
"""
        }
    }
    
    @classmethod
    def detect_template_type(cls, task_description: str, task_name: str) -> Optional[str]:
        """
        Detecta qual template usar baseado na descrição da tarefa
        
        Returns:
            Nome do template ou None se não detectar
        """
        desc_lower = task_description.lower()
        name_lower = task_name.lower()
        
        # Detectar fileira/linha
        if any(word in desc_lower or word in name_lower 
               for word in ['row', 'fileira', 'linha', 'line', 'fila']):
            return 'row'
        
        # Detectar empilhamento simples
        if any(word in desc_lower or word in name_lower 
               for word in ['stack', 'empilhar', 'stacking']):
            if 'pyramid' not in desc_lower and 'pyramid' not in name_lower:
                return 'stack'
        
        # Detectar pirâmide
        if 'pyramid' in desc_lower or 'pyramid' in name_lower or 'pirâmide' in desc_lower:
            return 'pyramid'
        
        return None
    
    @classmethod
    def get_template(cls, template_type: str) -> Optional[Dict]:
        """Retorna template específico"""
        return cls.TEMPLATES.get(template_type)
    
    @classmethod
    def should_use_template(cls, task_description: str, task_name: str) -> bool:
        """Decide se deve usar template ou geração livre"""
        template_type = cls.detect_template_type(task_description, task_name)
        return template_type is not None


def suggest_template_usage(task_description: str, task_name: str) -> Optional[str]:
    """
    Sugere se deve usar template e qual
    
    Returns:
        Nome do template sugerido ou None
    """
    return CodeTemplate.detect_template_type(task_description, task_name)