yfzhoucs/TinyLanguageRobots

1import random2import copy3import numpy as np4 5 6class Initializer:7    def __init__(self, config, obj_num_low=3, obj_num_high=5):8        self.config = config9        self.obj_num_low = obj_num_low10        self.obj_num_high = obj_num_high11        self.available_actions = [12            'push_forward', 13            'push_backward', 14            'push_left', 15            'push_right', 16            'rotate_clock', 17            'rotate_counterclock']18 19        self.verb_template = {20            'push_forward': [21                'push',22                'drag',23                'move',24                'get',25            ], 26            'push_backward': [27                'push',28                'drag',29                'move',30                'get',31            ], 32            'push_left': [33                'push',34                'drag',35                'move',36                'get',37            ],  38            'push_right': [39                'push',40                'drag',41                'move',42                'get',43            ],  44            'rotate_clock': [45                'rotate',46                'revolve',47                'turn',48                'spin',49            ],  50            'rotate_counterclock': [51                'rotate',52                'revolve',53                'turn',54                'spin',55            ], 56        }57        self.adv_template = {58            'push_forward': [59                'forward',60                'to the front',61                'ahead'62            ], 63            'push_backward': [64                'backward',65                'back',66                'to the back',67            ], 68            'push_left': [69                'to the left',70                'left',71                'to the left hand side',72            ],  73            'push_right': [74                'to the right',75                'right',76                'to the right hand side',77            ],  78            'rotate_clock': [79                'clockwise',80                'clock wise',81                'right'82            ],  83            'rotate_counterclock': [84                'counterclockwise',85                'anticlockwise',86                'anti clock wise',87                'counter clock wise',88                'left',89            ], 90        }91        self.np_template = {92            'orange': [93                'orange',94                'citrus',95                'sweet orange',96                'lime'97            ],98            'apple': [99                'apple',100                'red apple',101                'red delicious apple',102                'gala'103            ],104            'tomato':[105                'tomato'106            ],107            'strawberry':[108                'strawberry'109            ],110            'watermelon':[111                'watermelon'112            ],113            'banana':[114                'banana'115            ],116            'milk_bottle':[117                'bottle',118                'glass bottle',119                'milk bottle'120            ],121            'clock':[122                'clock',123                'timer',124                'watch'125            ],126            'camera':[127                'camera',128                'DSLR',129                'nikon',130                'canon'131            ]132        }133        return134    135    def _random_place_(self, positions):136        137        def l2(x1, y1, x2, y2):138            return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** (1/2)139 140        def not_too_close(positions, x, y):141            for pos in positions:142                if l2(x, y, pos['x'], pos['y']) <= 100:143                    return False144            return True145 146        done = False147 148        while not done:149            # https://programming.guide/random-point-within-circle.html150            a = random.random() * np.pi151            r = 450 * np.sqrt(random.random())152            x = r * np.cos(a) + self.config['desk_width'] * self.config['scale']153            y = r * np.sin(a) + 100154 155            if not_too_close(positions, x, y) and r > 250:156                done = True157        158        return {159            'x': x,160            'y': y,161            'z': 0162        }163 164    165    def get_config_and_task(self):166        available_objs = []167        for obj in self.config['objects']:168            if 'position' in self.config['objects'][obj]:169                available_objs.append(obj)170 171        num_obj = random.randint(self.obj_num_low, self.obj_num_high)172        sampled_objs = random.sample(available_objs, num_obj)173 174        config = {}175        for key in self.config:176            if not key == 'objects':177                config[key] = copy.deepcopy(self.config[key])178        179        config['objects'] = {}180        positions = []181        for obj in self.config['objects']:182            if 'position' not in self.config['objects'][obj]:183                config['objects'][obj] = copy.deepcopy(self.config['objects'][obj])184            elif obj in sampled_objs:185                obj_dict = copy.deepcopy(self.config['objects'][obj])186                obj_dict['position'] = self._random_place_(positions)187                positions.append(obj_dict['position'])188                config['objects'][obj] = obj_dict189 190        task = {191            'action': random.sample(self.available_actions, 1)[0],192            'target': random.sample(sampled_objs, 1)[0],193        }194 195        config['init_joints'] = np.random.uniform(-np.pi / 2, np.pi / 2, size=(4,))196        config['init_joints'][-1] *= 0.2197 198        self.new_config = config199        self.task = task200        print(config)201        print(task)202        return config, task203    204    def get_verb(self, verb):205        return random.sample(self.verb_template[verb], 1)[0]206    207    def get_adv(self, verb):208        return random.sample(self.adv_template[verb], 1)[0]209 210    def get_np(self, noun):211        return random.sample(self.np_template[noun], 1)[0]212 213    def get_sentence(self):214        v = self.get_verb(self.task['action'])215        adv = self.get_adv(self.task['action'])216        np = self.get_np(self.task['target'])217        sentence = v + ' ' + np + ' ' + adv218        return sentence