├── .gitignore ├── README.md ├── configuration_files ├── configs_barrier.json ├── configs_barrier_with_door.json ├── configs_pit.json ├── configs_single_goal_cube.json ├── configs_single_goal_cube_part_2.json ├── configs_single_goal_ramp.json ├── configurations_3_4_cube.json ├── configurations_3_4_ramp.json └── sample_single_goal.py ├── main.py ├── pre_proccess_data.py ├── prepare-dataset.py └── util ├── barrier_with_door.py ├── camera.py ├── create_scene.py ├── generate_configuration.py ├── imgs ├── fig1.png └── icon.gif ├── jump.py ├── obstacle.py ├── occuluder.py ├── onegoalscene.py ├── physics_agent.py ├── pit.py ├── ramp.py ├── replay_agent.py ├── replay_scene.py ├── sample_trajectory.py ├── scene_configuration.py ├── target_object.py ├── trajectory_generation.py ├── twogoalscene.py └── utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | data/* 2 | int_phys/* 3 | __pycache__/* 4 | .idea/* 5 | .DS_Store 6 | test.py 7 | test2.py 8 | 9 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AGENT: A Benchmark for Core Psychological Reasoning 2 | 3 | Code for procedurally synthesizing videos used for constructing the trials in the dataset, AGENT, proposed in the paper: [*AGENT: A Benchmark for Core Psychological Reasoning*](https://www.tshu.io/AGENT/AGENT.pdf). 4 | 5 | ![figure](util/imgs/icon.gif "AGENT video example.") 6 | 7 | 8 | ## Installation and Running 9 | * Clone the repository 10 | * Download agent 3d models [here](https://agent-dataset-storage.s3.us-east.cloud-object-storage.appdomain.cloud/agent_3d_models.zip) 11 | and extract the folder inside the `AGENT/` folder 12 | * You can run the script in two modes: 13 | * To run trajectories 14 | ``` 15 | python main.py --out-dir data --config-file configuration_files.json --config-start 0 --config-end 1 --scene-type scene_1 16 | ``` 17 | * To run replay 18 | ``` 19 | python main.py --out-dir data --config-file configuration_files.json --config-start 0 --config-end 100 --scene-type replay_scene 20 | ``` 21 | 22 | 23 | ## Dataset Structure Overview 24 | * We provide dataset in two bundles: Training & validation set, and Testing set 25 | * The dataset is divided into 4 scenarios. 26 | * scenario_1_goal_preferences 27 | * scenario_2_action_efficiency 28 | * scenario_3_unobserved_constraints 29 | * scenario_4_cost_reward_trade_offs 30 | * Each scenario is further divided into subtypes. 31 | * Each subtype then contains will then contain train set and validation set or testing set depending on which 32 | zip folder you are inside. 33 | * Each of the train, test and validation folders will contain the final trial sets. The number of trials is 34 | summarized in this ![figure](util/imgs/fig1.png "The number of trials for each type in AGENT.") 35 | 36 | ## Trial Structure 37 | * For each trial we provide RGB and Depth map images as well as state data 38 | * Each Trial folder is organized as : 39 | * Familiarization and expected or unexpected video or both depending on which set (train, validation, test) 40 | you are in and have similar structure. 41 | * Each folder contains: 42 | * images_c/ folder: Collection of numbered images/frames. Prefix used are: `img_`: RGB, `id_`: Segmentation 43 | mask, `depth_`: Depth map 44 | * Videos: `scene_c.mp4` RGB image video, `mask_scene_c.mp4` segmentation map video, `depth_scene_c.mp4` 45 | depth map video 46 | * state_info.json : This JSON contains list of states. The len of list is equal to number of RBG image/ 47 | Depth images/Segmentation mask. Each element in list is dictionary with keys being object name and value 48 | being object attributes like position, velocity, rotation, etc. [More information](#more-on-state-data) 49 | * scene_config.json: This contains dictionary of objects in the scene. This is used by the procedural generation 50 | script to recreate the scene. 51 | 52 | ## More on State Data 53 | * The state data which is list of dictionary, each of which represents state at each frame, uses different keys to 54 | represent objects. 55 | * These keys include: 56 | * `agent`: Agent in the scene. Scenes will always have 1 agent 57 | * `goal`, `goal_`: Goal objects in scene. If there are more than one goal object than it will have number 58 | indicating which side it is on. 1 is left and 2 is right of agent 59 | * `obstacle`, `obstacle_`: This object is the plain wall barrier. Scene with multiple barrier will have 60 | side numbers 61 | * `ramp_slope_`, `ramp_platform_`, `ramp_slope`, `ramp_platform`: Each ramp is composed of 62 | the platform and slope. Each piece has its own set of attributes 63 | * `barrier__with_door_`, `barrier_with_door_`: This is a plain wall with 64 | door cutout in the middle. This obstacle is composed of 3 pieces indicated by the piece number 65 | * `pit_bridge_`, `pit_side_part_`: Pit is formed by 2-3 pieces. In some scenario the pit 66 | sides will be connected by bridge 67 | * `occluder`: In Unobserved constraints' scenario the occluder is used to hide the obstacle. 68 | 69 | 70 | ## Cite 71 | If you use this code in your research, please consider citing the following the paper. 72 | 73 | ``` 74 | @inproceedings{shu2021agent, 75 | title={{AGENT}: A Benchmark for Core Psychological Reasoning}, 76 | author={Shu, Tianmin and Bhandwaldar, Abhishek and Gan, Chuang and Smith, Kevin A and Liu, Shari and Gutfreund, Dan and Spelke, Elizabeth and Tenenbaum, Joshua B and Ullman, Tomer D}, 77 | booktitle={The Thirty-eighth International Conference on Machine Learning (ICML)}, 78 | year={2021} 79 | } 80 | ``` 81 | -------------------------------------------------------------------------------- /configuration_files/sample_single_goal.py: -------------------------------------------------------------------------------- 1 | import random 2 | import json 3 | import itertools 4 | 5 | barrier_type = 'cube' 6 | NUM_AGENT_SHAPE = 1 7 | NUM_OBJECT_SHAPE = 1 8 | NUM_POS_X = 3 9 | NUM_POS_Z = 3 10 | NUM_OBSTACLE_HEIGHT = 4 11 | NUM_OBSTACLE_WIDTH = 2 12 | NUM_OBSTACLE_DEPTH = 4 13 | 14 | configs = [] 15 | 16 | valid = [] 17 | 18 | efficient = [] 19 | valid_efficient = [] 20 | 21 | for agent_shape in range(NUM_AGENT_SHAPE): 22 | for obj_shape in range(NUM_OBJECT_SHAPE): 23 | for agent_pos_x in range(NUM_POS_X): 24 | for obj_pos_x in range(NUM_POS_X): 25 | if agent_pos_x == obj_pos_x: continue 26 | for obstacle_pos_x in range(NUM_POS_X): 27 | if agent_pos_x == obstacle_pos_x or obj_pos_x == obstacle_pos_x: continue 28 | pos_z_lists = [list(range(NUM_POS_Z)), list(range(NUM_POS_Z)), list(range(NUM_POS_Z))] 29 | all_pos_z = list(itertools.product(*pos_z_lists)) 30 | for pos_z in all_pos_z: 31 | agent_pos_z, obj_pos_z, obstacle_pos_z = pos_z[0], pos_z[1], pos_z[2] 32 | for obstacle_height in range(NUM_OBSTACLE_HEIGHT): 33 | for obstacle_width in range(NUM_OBSTACLE_WIDTH): 34 | if obstacle_height == 0 and obstacle_width > 0: continue 35 | for obstacle_depth in range(NUM_OBSTACLE_DEPTH): 36 | config = {'barrier_type': barrier_type, 37 | 'id': len(configs), 38 | 'agent_shape': agent_shape, 39 | 'obj_shape': obj_shape, 40 | 'agent_pos_x': agent_pos_x, 41 | 'agent_pos_z': agent_pos_z, 42 | 'obj_pos_x': obj_pos_x, 43 | 'obj_pos_z': obj_pos_z, 44 | 'obstacle_pos_x': obstacle_pos_x, 45 | 'obstacle_pos_z': obstacle_pos_z, 46 | 'obstacle_height': obstacle_height, 47 | 'obstacle_width': obstacle_width, 48 | 'obstacle_depth': obstacle_depth} 49 | configs.append(config) 50 | 51 | 52 | print(len(configs)) 53 | 54 | with open('configs_single_goal_{}.json'.format(barrier_type), 'w') as outfile: 55 | json.dump(configs, outfile, indent=4) 56 | 57 | 58 | # barrier_type = 'ramp' 59 | # NUM_AGENT_SHAPE = 1 60 | # NUM_OBJECT_SHAPE = 1 61 | # NUM_POS_X = 3 62 | # NUM_POS_Z = 3 63 | # NUM_RAMP_HEIGHT = 2 64 | # NUM_RAMP_ROTATION = 4 65 | # NUM_OBJ_POS_RAMP = 2 66 | 67 | # configs = [] 68 | 69 | # valid = [] 70 | 71 | # efficient = [] 72 | # valid_efficient = [] 73 | 74 | # for agent_shape in range(NUM_AGENT_SHAPE): 75 | # for obj_shape in range(NUM_OBJECT_SHAPE): 76 | # for agent_pos_x in range(NUM_POS_X): 77 | # for obstacle_pos_x in range(NUM_POS_X): 78 | # if agent_pos_x == obstacle_pos_x: continue 79 | # pos_z_lists = [list(range(NUM_POS_Z)), list(range(NUM_POS_Z))] 80 | # all_pos_z = list(itertools.product(*pos_z_lists)) 81 | # for pos_z in all_pos_z: 82 | # agent_pos_z, obstacle_pos_z = pos_z[0], pos_z[1] 83 | # for ramp_height in range(NUM_RAMP_HEIGHT): 84 | # for ramp_rotation in range(NUM_RAMP_ROTATION): 85 | # for obj_pos_ramp in range(NUM_OBJ_POS_RAMP): 86 | # config = {'barrier_type': barrier_type, 87 | # 'id': len(configs), 88 | # 'agent_shape': agent_shape, 89 | # 'obj_shape': obj_shape, 90 | # 'agent_pos_x': agent_pos_x, 91 | # 'agent_pos_z': agent_pos_z, 92 | # 'obstacle_pos_x': obstacle_pos_x, 93 | # 'obstacle_pos_z': obstacle_pos_z, 94 | # 'ramp_height': ramp_height, 95 | # 'ramp_rotation': ramp_rotation, 96 | # 'obj_pos_ramp': obj_pos_ramp} 97 | # configs.append(config) 98 | 99 | 100 | # print(len(configs)) 101 | 102 | # with open('configs_single_goal_{}.json'.format(barrier_type), 'w') as outfile: 103 | # json.dump(configs, outfile, indent=4) 104 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | from tdw.controller import Controller 2 | from util import utils 3 | import os 4 | import json 5 | import tqdm 6 | from util.onegoalscene import OneGoalScene 7 | from util.twogoalscene import TwoGoalScene 8 | from util.replay_scene import ReplayScene 9 | import requests 10 | import argparse 11 | import random 12 | 13 | v = -0.03 14 | h = 0.09 15 | 16 | 17 | class AGENTController(Controller): 18 | def run(self, args): 19 | self.start() 20 | # Create an empty room. 21 | if args.scene_type == "scene_1": 22 | camera_obj = utils.create_base_scene(self, args, single_camera=True) 23 | self.run_single_goal_scenario(camera_obj, args) 24 | elif args.scene_type == "scene_3_4": 25 | camera_obj = utils.create_base_scene(self, args, single_camera=True, scene_3_camera=True) 26 | self.run_two_goal_scenario(camera_obj, args) 27 | elif args.scene_type == "replay_scene": 28 | camera_obj = utils.create_base_scene(self, args, single_camera=True, scene_3_camera=True 29 | if args.replay_scene_type == "scene_3" else False) 30 | self.run_play_back(camera_obj, args) 31 | elif args.scene_type == "color_sampler": 32 | camera_obj = utils.create_base_scene(self, args, single_camera=True, scene_3_camera=True) 33 | utils.create_color_samples(self) 34 | self.communicate({"$type": "terminate"}) 35 | # Procedural 36 | 37 | def run_play_back(self, camera_obj, args): 38 | replay_scene = ReplayScene(self, output_dir="replay_data/config_20", camera_obj=camera_obj, config_path="data/config_20", args=args) 39 | if args.replay_scene_type == "scene_1": 40 | replay_scene.scenario_2_action_efficiency(args) 41 | elif args.replay_scene_type == "scene_2": 42 | replay_scene.scenario_3_unobserved_constraints(args) 43 | elif args.replay_scene_type == "scene_3": 44 | replay_scene.scenario_1_goal_preferences(args) 45 | else: 46 | replay_scene.scenario_4_cost_reward_trade_offs(args) 47 | 48 | def run_single_goal_scenario(self, camera_obj, args): 49 | scene_1 = OneGoalScene(self, camera_obj, args) 50 | data = [ 51 | # { 52 | # "barrier_type": "pit-with-bridge", 53 | # "id": 0, 54 | # "agent_shape": 0, 55 | # "obj_shape": 0, 56 | # "agent_pos_x": 0, 57 | # "agent_pos_z": 0, 58 | # "obj_pos_x": 2, 59 | # "obj_pos_z": 2, 60 | # "pit_width": 1, 61 | # "pit_depth": -1, 62 | # "bridge_z": 1, 63 | # "obstacle_pos_x": 1, 64 | # "obstacle_pos_z": 1, 65 | # }, 66 | # { 67 | # "barrier_type": "barrier_with_door", 68 | # "id": 3, 69 | # "agent_shape": 2, 70 | # "obj_shape": 2, 71 | # "agent_pos_x": 0, 72 | # "agent_pos_z": 1, 73 | # "obj_pos_x": 2, 74 | # "obj_pos_z": 1, 75 | # "obstacle_pos_x": 1, 76 | # "obstacle_pos_z": 1, 77 | # "obstacle_height": 6, 78 | # "obstacle_width": 0, 79 | # "obstacle_depth": 5, 80 | # }, 81 | # { 82 | # "barrier_type": "barrier_with_door", 83 | # "id": 4, 84 | # "agent_shape": 2, 85 | # "obj_shape": 0, 86 | # "agent_pos_x": 0, 87 | # "agent_pos_z": 1, 88 | # "obj_pos_x": 2, 89 | # "obj_pos_z": 1, 90 | # "obstacle_pos_x": 1, 91 | # "obstacle_pos_z": 1, 92 | # "obstacle_height": 6, 93 | # "obstacle_width": 0, 94 | # "obstacle_depth": 5, 95 | # }, 96 | # { 97 | # "barrier_type": "platform-cube", 98 | # "id": 6, 99 | # "agent_shape": 2, 100 | # "obj_shape": 0, 101 | # "agent_pos_x": 0, 102 | # "agent_pos_z": 2, 103 | # "obj_pos_x": 2, 104 | # "obj_pos_z": 0, 105 | # "obstacle_pos_x": 1, 106 | # "obstacle_pos_z": 0, 107 | # "obstacle_height": 0 108 | # }, 109 | { 110 | "barrier_type": "cube", 111 | "id": 7, 112 | "agent_shape": 2, 113 | "obj_shape": 0, 114 | "agent_pos_x": 2, 115 | "agent_pos_z": 2, 116 | "obj_pos_x": 0, 117 | "obj_pos_z": 2, 118 | "obstacle_pos_x": 1, 119 | "obstacle_pos_z": 2, 120 | "obstacle_height": 6, 121 | "obstacle_width": 0, 122 | "obstacle_depth": 4, 123 | }, 124 | # { 125 | # "barrier_type": "cube", 126 | # "id": 8, 127 | # "agent_shape": 2, 128 | # "obj_shape": 0, 129 | # "agent_pos_x": 0, 130 | # "agent_pos_z": 0, 131 | # "obj_pos_x": 2, 132 | # "obj_pos_z": 2, 133 | # "obstacle_pos_x": 1, 134 | # "obstacle_pos_z": 2, 135 | # "obstacle_height": 4, 136 | # "obstacle_width": 0, 137 | # "obstacle_depth": 2, 138 | # }, 139 | # { 140 | # "barrier_type": "platform", 141 | # "id": 19, 142 | # "agent_shape": 0, 143 | # "agent_pos_x": 0, 144 | # "agent_pos_z": 2, 145 | # "obj_shape": 0, 146 | # "obj_pos_ramp_x": 2, 147 | # "obj_pos_ramp_z": 2, 148 | # "ramp_height": 1, 149 | # }, 150 | # { 151 | # "barrier_type": "platform", 152 | # "id": 20, 153 | # "agent_shape": 0, 154 | # "agent_pos_x": 2, 155 | # "agent_pos_z": 2, 156 | # "obj_shape": 0, 157 | # "obj_pos_ramp_x": 0, 158 | # "obj_pos_ramp_z": 2, 159 | # "ramp_height": 1, 160 | # }, 161 | # { 162 | # "barrier_type": "platform", 163 | # "id": 21, 164 | # "agent_shape": 0, 165 | # "agent_pos_x": 0, 166 | # "agent_pos_z": 2, 167 | # "obj_shape": 0, 168 | # "obj_pos_ramp_x": 2, 169 | # "obj_pos_ramp_z": 2, 170 | # "ramp_height": 2, 171 | # }, 172 | # { 173 | # "barrier_type": "ramp", 174 | # "id": 22, 175 | # "agent_shape": 0, 176 | # "agent_pos_x": 0, 177 | # "agent_pos_z": 0, 178 | # "obj_shape": 0, 179 | # "obj_pos_ramp": 1, 180 | # "ramp_height": 1, 181 | # "ramp_rotation": 2 182 | # }, 183 | # { 184 | # "barrier_type": "ramp", 185 | # "id": 21, 186 | # "agent_shape": 0, 187 | # "agent_pos_x": 2, 188 | # "agent_pos_z": 1, 189 | # "obj_shape": 0, 190 | # "obj_pos_ramp": 0, 191 | # "ramp_height": 1, 192 | # "ramp_rotation": 0 193 | # }, 194 | # { 195 | # "barrier_type": "cube", 196 | # "id": 10, 197 | # "agent_shape": 0, 198 | # "obj_shape": 10, 199 | # "agent_pos_x": 0, 200 | # "agent_pos_z": 1, 201 | # "obj_pos_x": 2, 202 | # "obj_pos_z": 1, 203 | # "obstacle_pos_x": 1, 204 | # "obstacle_pos_z": 1, 205 | # "obstacle_height": 5, 206 | # "obstacle_width": 0, 207 | # "obstacle_depth": 3, 208 | # }, 209 | 210 | # { 211 | # "barrier_type": "ramp", 212 | # "id": 21, 213 | # "agent_shape": 1, 214 | # "obj_shape": 0, 215 | # "agent_pos_x": 0, 216 | # "agent_pos_z": 1, 217 | # "obstacle_pos_x": 1, 218 | # "obstacle_pos_z": 1, 219 | # "ramp_height": 1, 220 | # "ramp_rotation": 0, 221 | # "obj_pos_ramp": 1, 222 | # }, 223 | # # { 224 | # # "barrier_type": "ramp", 225 | # # "id": 21, 226 | # # "agent_shape": 1, 227 | # # "obj_shape": 0, 228 | # # "agent_pos_x": 0, 229 | # # "agent_pos_z": 1, 230 | # # "obstacle_pos_x": 1, 231 | # # "obstacle_pos_z": 1, 232 | # # "ramp_height": 0, 233 | # # "ramp_rotation": 1, 234 | # # "obj_pos_ramp": 1, 235 | # # }, 236 | # # # # { 237 | # # # # "barrier_type": "ramp", 238 | # # # # "id": 22, 239 | # # # # "agent_shape": 1, 240 | # # # # "obj_shape": 0, 241 | # # # # "agent_pos_x": 0, 242 | # # # # "agent_pos_z": 1, 243 | # # # # "obstacle_pos_x": 1, 244 | # # # # "obstacle_pos_z": 1, 245 | # # # # "ramp_height": 0, 246 | # # # # "ramp_rotation": 2, 247 | # # # # "obj_pos_ramp": 1, 248 | # # # # "obstacle_height": 0, 249 | # # # # "obstacle_width": 0, 250 | # # # # "obstacle_depth": 2, 251 | # # # # "approach": "jump_ramp_1" 252 | # # # # }, 253 | # { 254 | # "barrier_type": "ramp", 255 | # "id": 23, 256 | # "agent_shape": 1, 257 | # "obj_shape": 0, 258 | # "agent_pos_x": 0, 259 | # "agent_pos_z": 1, 260 | # "obstacle_pos_x": 1, 261 | # "obstacle_pos_z": 1, 262 | # "ramp_height": 2, 263 | # "ramp_rotation": 2, 264 | # "obj_pos_ramp": 1, 265 | # }, 266 | # { 267 | # "barrier_type": "ramp", 268 | # "id": 26, 269 | # "agent_shape": 1, 270 | # "obj_shape": 0, 271 | # "agent_pos_x": 0, 272 | # "agent_pos_z": 1, 273 | # "obstacle_pos_x": 1, 274 | # "obstacle_pos_z": 1, 275 | # "ramp_height": 1, 276 | # "ramp_rotation": 0, 277 | # "obj_pos_ramp": 1, 278 | # } 279 | # # # 280 | # # # # { 281 | # # # # "barrier_type": "cube", 282 | # # # # "id": 13, 283 | # # # # "agent_shape": 0, 284 | # # # # "obj_shape": 0, 285 | # # # # "agent_pos_x": 0, 286 | # # # # "agent_pos_z": 1, 287 | # # # # "obj_pos_x": 3, 288 | # # # # "obj_pos_z": 0, 289 | # # # # "obstacle_pos_x": 1, 290 | # # # # "obstacle_pos_z": 0, 291 | # # # # "obstacle_height": 3, 292 | # # # # "obstacle_width": 0, 293 | # # # # "obstacle_depth": 2, 294 | # # # # "jump_height": 0, 295 | # # # # "go_around": "right" 296 | # # # # }, 297 | # # # # { 298 | # # # # "barrier_type": "cube", 299 | # # # # "id": 14, 300 | # # # # "agent_shape": 0, 301 | # # # # "obj_shape": 0, 302 | # # # # "agent_pos_x": 0, 303 | # # # # "agent_pos_z": 1, 304 | # # # # "obj_pos_x": 3, 305 | # # # # "obj_pos_z": 0, 306 | # # # # "obstacle_pos_x": 1, 307 | # # # # "obstacle_pos_z": 0, 308 | # # # # "obstacle_height": 3, 309 | # # # # "obstacle_width": 0, 310 | # # # # "obstacle_depth": 2, 311 | # # # # "jump_height": 3, 312 | # # # # }, 313 | # # # # { 314 | # # # # "barrier_type": "cube", 315 | # # # # "id": 15, 316 | # # # # "agent_shape": 0, 317 | # # # # "obj_shape": 0, 318 | # # # # "agent_pos_x": 0, 319 | # # # # "agent_pos_z": 1, 320 | # # # # "obj_pos_x": 3, 321 | # # # # "obj_pos_z": 0, 322 | # # # # "obstacle_pos_x": 1, 323 | # # # # "obstacle_pos_z": 0, 324 | # # # # "obstacle_height": 3, 325 | # # # # "obstacle_width": 0, 326 | # # # # "obstacle_depth": 2, 327 | # # # # "jump_height": 4, 328 | # # # # }, 329 | # # # # { 330 | # # # # "barrier_type": "cube", 331 | # # # # "id": 16, 332 | # # # # "agent_shape": 0, 333 | # # # # "obj_shape": 0, 334 | # # # # "agent_pos_x": 0, 335 | # # # # "agent_pos_z": 1, 336 | # # # # "obj_pos_x": 3, 337 | # # # # "obj_pos_z": 0, 338 | # # # # "obstacle_pos_x": 1, 339 | # # # # "obstacle_pos_z": 0, 340 | # # # # "obstacle_height": 3, 341 | # # # # "obstacle_width": 0, 342 | # # # # "obstacle_depth": 2, 343 | # # # # "jump_height": 5, 344 | # # # # }, 345 | # # # 346 | # { 347 | # "barrier_type": "pit", 348 | # "id": 0, 349 | # "agent_shape": 0, 350 | # "obj_shape": 0, 351 | # "agent_pos_x": 0, 352 | # "agent_pos_z": 1, 353 | # "obj_pos_x": 2, 354 | # "obj_pos_z": 0, 355 | # "pit_width": 0, 356 | # "pit_depth": -1, 357 | # "obstacle_pos_x": 1, 358 | # "obstacle_pos_z": 1, 359 | # }, 360 | # { 361 | # "barrier_type": "pit-with-bridge", 362 | # "id": 1, 363 | # "agent_shape": 0, 364 | # "obj_shape": 0, 365 | # "agent_pos_x": 0, 366 | # "agent_pos_z": 1, 367 | # "obj_pos_x": 2, 368 | # "obj_pos_z": 1, 369 | # "pit_width": 1, 370 | # "bridge_z": 1, 371 | # "pit_depth": -1, 372 | # "obstacle_pos_x": 1, 373 | # "obstacle_pos_z": 1 374 | # }, 375 | # # # { 376 | # # # "barrier_type": "pit", 377 | # # # "id": 18, 378 | # # # "agent_shape": 0, 379 | # # # "obj_shape": 0, 380 | # # # "agent_pos_x": 0, 381 | # # # "agent_pos_z": 1, 382 | # # # "obj_pos_x": 3, 383 | # # # "obj_pos_z": 1, 384 | # # # "pit_width": 0, 385 | # # # "jump_width": 1, 386 | # # # }, 387 | # # # # { 388 | # # # # "barrier_type": "pit", 389 | # # # # "id": 19, 390 | # # # # "agent_shape": 0, 391 | # # # # "obj_shape": 0, 392 | # # # # "agent_pos_x": 0, 393 | # # # # "agent_pos_z": 1, 394 | # # # # "obj_pos_x": 3, 395 | # # # # "obj_pos_z": 1, 396 | # # # # "pit_width": 0, 397 | # # # # "jump_width": 2, 398 | # # # # }, 399 | # # # 400 | # # # 401 | # # # 402 | ] 403 | # width -> code 404 | # 6 -> 0 405 | # 5.78 -> 1 406 | # 5.5 -> 2 407 | # 5.25 -> 3 408 | # 5 -> 4 409 | # pos_z -> code 410 | # -2.795 -> 0 411 | # -2.4 -> 1 412 | # -1.917 -> 2 413 | # 414 | # Pos_x -> code 415 | # -1.32 -> 2 416 | # 1.08 -> 0 417 | # with open(args.config_file, "r") as fp: 418 | # data = json.load(fp) 419 | 420 | for i in tqdm.tqdm(range(args.config_start, args.config_end, 1)): 421 | e = data[i] 422 | e["id"] += 1 423 | file_id = e["id"] 424 | e["dir_name"] = os.path.join(args.out_dir, f"config_{file_id}") 425 | scene_1.execute_config(config_file=e, output_dir=args.out_dir) 426 | 427 | def run_two_goal_scenario(self, camera_obj, args): 428 | scene_3_4 = TwoGoalScene(self, camera_obj, args) 429 | data = [ 430 | { 431 | "barrier_type": "ramp", 432 | "id": 1, 433 | "agent_shape": 0, 434 | "obj_shape": 0, 435 | "agent_pos_z": 0, 436 | "ramp_height_1": 0, 437 | "ramp_rotation_1": 0, 438 | "ramp_height_2": 1, 439 | "ramp_rotation_2": 2 440 | } 441 | # { 442 | # "barrier_type": "pit-with-bridge", 443 | # "id": 0, 444 | # "agent_shape": 0, 445 | # "obj_shape": 0, 446 | # "agent_pos_x": 0, 447 | # "agent_pos_z": 1, 448 | # "obj_1_pos_x": 0, 449 | # "obj_1_pos_z": 1, 450 | # "obj_shape_1": 1, 451 | # "obj_shape_2": 1, 452 | # "obj_2_pos_x": 4, 453 | # "obj_2_pos_z": 1, 454 | # "pit_width_1": 2, 455 | # "pit_depth_1": 2, 456 | # "pit_width_2": 0, 457 | # "pit_depth_2": 2, 458 | # "bridge_1_z": 1, 459 | # "bridge_2_z": 1, 460 | # "obstacle_1_pos_x": 1, 461 | # "obstacle_1_pos_z": 1, 462 | # "obstacle_2_pos_x": 3, 463 | # "obstacle_2_pos_z": 1, 464 | # }, 465 | # { 466 | # "barrier_type": "pit", 467 | # "id": 0, 468 | # "agent_shape": 0, 469 | # "obj_shape": 0, 470 | # "agent_pos_x": 0, 471 | # "agent_pos_z": 2, 472 | # "obj_1_pos_x": 0, 473 | # "obj_1_pos_z": 1, 474 | # "obj_2_pos_x": 4, 475 | # "obj_2_pos_z": 1, 476 | # "pit_width_1": 0, 477 | # "pit_depth_1": 2, 478 | # "pit_width_2": -1, 479 | # "pit_depth_2": 2, 480 | # "obstacle_1_pos_x": 1, 481 | # "obstacle_1_pos_z": 1, 482 | # "obstacle_2_pos_x": 3, 483 | # "obstacle_2_pos_z": 1, 484 | # }, 485 | # { 486 | # "barrier_type": "platform", 487 | # "id": 1, 488 | # "agent_shape": 0, 489 | # "obj_shape_1": 1, 490 | # "obj_shape_2": 3, 491 | # "obj_shape": 0, 492 | # "agent_pos_z": 1, 493 | # "ramp_height_1": 1, 494 | # "ramp_rotation_1": 2, 495 | # "ramp_height_2": 1, 496 | # "ramp_rotation_2": 0 497 | # } 498 | # { 499 | # "barrier_type": "platform", 500 | # "id": 0, 501 | # "agent_shape": 0, 502 | # "obj_shape": 0, 503 | # "agent_pos_z": 1, 504 | # "obj_pos_ramp_x_1": 0, 505 | # "obj_pos_ramp_z_1": 2, 506 | # "obj_pos_ramp_x_2": 0, 507 | # "obj_pos_ramp_z_2": 2, 508 | # "ramp_height_1": -1, 509 | # "ramp_height_2": 1, 510 | # } 511 | # { 512 | # "barrier_type": "platform", 513 | # "barrier_type_1": "barrier_with_door", 514 | # "id": 1, 515 | # "agent_shape": 0, 516 | # "obj_shape": 0, 517 | # "agent_pos_z": 1, 518 | # "ramp_height_1": -1, 519 | # "ramp_rotation_1": 0, 520 | # "ramp_height_2": 1, 521 | # "ramp_rotation_2": 3, 522 | # "obstacle_1_pos_x": 1, 523 | # "obstacle_1_pos_z": 1, 524 | # "obstacle_1_height": 6, 525 | # "obstacle_1_width": 0, 526 | # "obstacle_1_depth": 3, 527 | # # "obj_2_pos_x": 4, 528 | # # "obj_2_pos_z": 1, 529 | # }, 530 | # { 531 | # "id": 2, 532 | # "barrier_type": "cube", 533 | # "agent_shape": 0, 534 | # "obj_shape_1": 1, 535 | # "obj_shape_2": 1, 536 | # "agent_pos_z": 1, 537 | # 538 | # "obj_1_pos_x": 0, 539 | # "obj_1_pos_z": 0, 540 | # "obj_2_pos_x": 4, 541 | # "obj_2_pos_z": 0, 542 | # 543 | # "obstacle_1_pos_x": 1, 544 | # "obstacle_1_pos_z": 1, 545 | # "obstacle_1_height": 1, 546 | # "obstacle_1_width": 0, 547 | # "obstacle_1_depth": 3, 548 | # 549 | # "obstacle_2_pos_x": 3, 550 | # "obstacle_2_pos_z": 1, 551 | # "obstacle_2_height": 1, 552 | # "obstacle_2_width": 0, 553 | # "obstacle_2_depth": 3, 554 | # }, 555 | # { 556 | # "id": 3, 557 | # "barrier_type": "barrier_with_door", 558 | # "agent_shape": 0, 559 | # "obj_shape": 0, 560 | # "agent_pos_z": 2, 561 | # 562 | # "obj_1_pos_x": 0, 563 | # "obj_1_pos_z": 2, 564 | # "obj_2_pos_x": 4, 565 | # "obj_2_pos_z": 2, 566 | # 567 | # "obstacle_1_pos_x": 1, 568 | # "obstacle_1_pos_z": 2, 569 | # "obstacle_1_height": 6, 570 | # "obstacle_1_width": 0, 571 | # "obstacle_1_depth": 3, 572 | # 573 | # "obstacle_2_pos_x": 3, 574 | # "obstacle_2_pos_z": 2, 575 | # "obstacle_2_height": 0, 576 | # "obstacle_2_width": 0, 577 | # "obstacle_2_depth": 3, 578 | # } 579 | # { 580 | # "id": 133, 581 | # "barrier_type": "cube", 582 | # "agent_shape": 0, 583 | # "obj_shape": 0, 584 | # "agent_pos_z": 1, 585 | # 586 | # "obj_1_pos_x": 0, 587 | # "obj_1_pos_z": 0, 588 | # "obj_2_pos_x": 4, 589 | # "obj_2_pos_z": 0, 590 | # 591 | # "obstacle_1_pos_x": 1, 592 | # "obstacle_1_pos_z": 2, 593 | # "obstacle_1_height": 6, 594 | # "obstacle_1_width": 3, 595 | # "obstacle_1_depth": 4, 596 | # 597 | # "obstacle_2_pos_x": 3, 598 | # "obstacle_2_pos_z": 2, 599 | # "obstacle_2_height": 6, 600 | # "obstacle_2_width": 3, 601 | # "obstacle_2_depth": 4, 602 | # 603 | # } 604 | # { 605 | # "id": 14, 606 | # "barrier_type": "cube", 607 | # "agent_shape": 0, 608 | # "obj_shape": 0, 609 | # "agent_pos_z": 2, 610 | # 611 | # "obj_1_pos_x": 0, 612 | # "obj_1_pos_z": 1, 613 | # "obj_2_pos_x": 4, 614 | # "obj_2_pos_z": 1, 615 | # 616 | # "obstacle_1_pos_x": 1, 617 | # "obstacle_1_pos_z": 2, 618 | # "obstacle_1_height": 5, 619 | # "obstacle_1_width": 3, 620 | # "obstacle_1_depth": 4, 621 | # 622 | # "obstacle_2_pos_x": 3, 623 | # "obstacle_2_pos_z": 2, 624 | # "obstacle_2_height": 3, 625 | # "obstacle_2_width": 1, 626 | # "obstacle_2_depth": 4, 627 | # 628 | # } 629 | ] 630 | # with open(args.config_file, "r") as fp: 631 | # data = json.load(fp) 632 | for i in tqdm.tqdm(range(args.config_start, args.config_end, 1)): 633 | e = data[i] 634 | file_id = e["id"] 635 | e["dir_name"] = os.path. join(args.out_dir, f"config_{file_id}") 636 | scene_3_4.execute_config(config_file=e, output_dir=args.out_dir) 637 | 638 | 639 | def create_tdw(port): 640 | url = "http://localhost:5000/get_tdw" 641 | data = { 642 | 'ip_address': "localhost", 643 | 'port': port 644 | } 645 | response = requests.post(url, json=json.dumps(data)) 646 | print(response.status_code, response.reason) 647 | docker_id = response.json()['docker_id'] 648 | return docker_id 649 | 650 | 651 | def kill_tdw(docker_id): 652 | url = "http://localhost:5000/kill_tdw" 653 | data = { 654 | "container_id": docker_id 655 | } 656 | response = requests.post(url, json=json.dumps(data)) 657 | print(response.status_code, response.reason) 658 | 659 | 660 | if __name__ == "__main__": 661 | parser = argparse.ArgumentParser(description='ProcGen') 662 | parser.add_argument( 663 | '--out-dir', 664 | default="data", 665 | help='Output directory') 666 | parser.add_argument( 667 | '--config-file', 668 | default="configuration.json", 669 | help='Configuration file for scene generation') 670 | parser.add_argument( 671 | '--config-start', 672 | default=0, 673 | type=int, 674 | help='Configuration file for scene generation') 675 | parser.add_argument( 676 | '--config-end', 677 | default=1, 678 | type=int, 679 | help='Configuration file for scene generation') 680 | parser.add_argument( 681 | '--port', 682 | default=1071, 683 | type=int, 684 | help='Configuration file for scene generation') 685 | parser.add_argument( 686 | '--set-num', 687 | default=1, 688 | type=int, 689 | help='Number of sets') 690 | parser.add_argument( 691 | '--screen-scale', 692 | default=80, 693 | type=int, 694 | help='Size of display') 695 | parser.add_argument( 696 | '--random-seed', 697 | default="saf203jdnflk", 698 | type=str, 699 | help='Random seed') 700 | parser.add_argument( 701 | '--input-dir', 702 | default="data", 703 | help='Input directory path for replay agent') 704 | parser.add_argument( 705 | '--scene-type', 706 | default="scene_1", 707 | help='Input directory path for replay agent') 708 | parser.add_argument( 709 | '--replay-scene-type', 710 | default="scene_1", 711 | help='Replay scene type') 712 | parser.add_argument( 713 | '--test-type', 714 | default="Type_1_0", 715 | help='Replay scene type') 716 | parser.add_argument( 717 | '--remote', 718 | action='store_true', 719 | default=False, 720 | help='Input directory path for replay agent') 721 | parser.add_argument( 722 | '--generate-video', 723 | action='store_true', 724 | default=False, 725 | help='Should video be generated?') 726 | parser.add_argument( 727 | '--use-novel-objects', 728 | action='store_true', 729 | default=False, 730 | help='Should novel objects be used ?') 731 | parser.add_argument( 732 | '--enable-image', 733 | action='store_true', 734 | default=False, 735 | help='Should novel objects be used ?') 736 | parser.add_argument( 737 | '--clear-images', 738 | action='store_true', 739 | default=False, 740 | help='Should novel objects be used ?') 741 | parser.add_argument( 742 | '--tmp', 743 | action='store_true', 744 | default=False, 745 | help='Tmp ?') 746 | parser.add_argument( 747 | '--reconstruct-3d', 748 | action='store_true', 749 | default=False, 750 | help='Reconstruct trajectory from data') 751 | parser.add_argument( 752 | '--follow-camera', 753 | action='store_true', 754 | default=False, 755 | help='Enable first person perspective') 756 | parser.add_argument( 757 | '--version', 758 | default="v2", 759 | help='Replay scene type') 760 | args = parser.parse_args() 761 | random.seed(args.random_seed) 762 | docker_id = None 763 | if args.scene_type == "replay_scene": 764 | args.config_end = args.config_start + args.set_num 765 | config_start = args.config_start 766 | config_end = args.config_end 767 | started = False 768 | for start in range(config_start, config_end, 50): 769 | args.config_end = start + 50 if start + 50 <= config_end else config_end 770 | args.config_start = start 771 | args.set_num = args.config_end - args.config_start 772 | if args.remote: 773 | AGENTController(port=args.port, launch_build=False).run(args) 774 | else: 775 | if not started: 776 | started = True 777 | AGENTController(port=args.port, launch_build=True).run(args) -------------------------------------------------------------------------------- /prepare-dataset.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | from os import listdir, makedirs 3 | import random 4 | from shutil import copytree, rmtree, copyfile 5 | import os 6 | from os.path import basename 7 | import json 8 | from pre_proccess_data import process_trajectory_state_data 9 | import pickle 10 | import shutil 11 | 12 | def save_fam_video(fam_paths, target_path): 13 | for e in fam_paths: 14 | if os.path.isdir(join(target_path, basename(e))): 15 | rmtree(join(target_path, basename(e))) 16 | file_name = basename(e) 17 | file_name = file_name.replace("amilarization", "amiliarization") 18 | copytree(e, join(target_path, file_name)) 19 | 20 | 21 | 22 | def create_dataset(base_paths, target_base_path, test_only, size=[0], versions=["version 1"], surprise_expected_reversed=False): 23 | if not test_only: 24 | create_train_data(base_paths, target_base_path, size, versions, surprise_expected_reversed) 25 | else: 26 | create_test_data(base_paths, target_base_path, versions, surprise_expected_reversed) 27 | 28 | 29 | def create_train_data(base_paths, target_base_path, sizes, versions, surprise_expected_reversed=False): 30 | train_set_no = 0 31 | val_set_no = 0 32 | train_path, val_path = join(target_base_path, "train_set"), join(target_base_path, "validation_set") 33 | makedirs(train_path, exist_ok=True) 34 | makedirs(val_path, exist_ok=True) 35 | # makedirs(test_path, exist_ok=True) 36 | for base_path, size, version in zip(base_paths, sizes, versions): 37 | trials = [join(base_path, "Agent_0", e) for e in listdir(join(base_path, "Agent_0"))] 38 | # Temp 39 | random.shuffle(trials) 40 | trials = trials[:size] 41 | val_list, train_list = trials[0:round(0.2*len(trials))], trials[round(0.2*len(trials)):] 42 | for trial in trials: 43 | sub_dirs = listdir(join(trial)) 44 | fam_dir = [join(trial, e) for e in sub_dirs if "fam" in e.lower()] 45 | test_a = [join(trial, e) for e in sub_dirs if "_A" in e][0] 46 | test_b = [join(trial, e) for e in sub_dirs if "_B" in e][0] 47 | if not surprise_expected_reversed: 48 | surprising, expected = test_a, test_b 49 | else: 50 | surprising, expected = test_b, test_a 51 | if trial in train_list: 52 | if os.path.isdir(join(train_path, f"Trial_{train_set_no}")): 53 | rmtree(join(train_path, f"Trial_{train_set_no}")) 54 | makedirs(join(train_path, f"Trial_{train_set_no}")) 55 | save_fam_video(fam_dir, join(train_path, f"Trial_{train_set_no}")) 56 | copytree(join(trial, expected), 57 | join(train_path, f"Trial_{train_set_no}", "expected")) 58 | open(join(train_path, f"Trial_{train_set_no}", version), "w") 59 | train_set_no += 1 60 | else: 61 | if os.path.isdir(join(val_path, f"Trial_{val_set_no}")): 62 | rmtree(join(val_path, f"Trial_{val_set_no}")) 63 | makedirs(join(val_path, f"Trial_{val_set_no}"), exist_ok=True) 64 | save_fam_video(fam_dir, join(val_path, f"Trial_{val_set_no}")) 65 | copytree(join(trial, expected), 66 | join(val_path, f"Trial_{val_set_no}", "expected")) 67 | open(join(val_path, f"Trial_{val_set_no}", version), "w") 68 | val_set_no += 1 69 | # if os.path.isdir(join(target_base_path, "test_set", f"Set_{test_set_no}")): 70 | # rmtree(join(target_base_path, "test_set", f"Set_{test_set_no}")) 71 | # makedirs(join(target_base_path, "test_set", f"Set_{test_set_no}"), exist_ok=True) 72 | # save_fam_video(fam_dir, join(target_base_path, "test_set", f"Set_{test_set_no}")) 73 | # copytree(join(trial, surprising), join(target_base_path, "test_set", f"Set_{test_set_no}", "surprising")) 74 | # test_set_no += 1 75 | # if os.path.isdir(join(target_base_path, "test_set", f"Set_{test_set_no}")): 76 | # rmtree(join(target_base_path, "test_set", f"Set_{test_set_no}")) 77 | # makedirs(join(target_base_path, "test_set", f"Set_{test_set_no}"), exist_ok=True) 78 | # save_fam_video(fam_dir, join(target_base_path, "test_set", f"Set_{test_set_no}")) 79 | # copytree(join(trial, expected), join(target_base_path, "test_set", f"Set_{test_set_no}", "expected")) 80 | # test_set_no += 1 81 | 82 | 83 | def create_test_data(base_paths, target_base_path, versions, surprise_expected_reversed): 84 | test_set_no = 0 85 | test_path = join(target_base_path, "test_set") 86 | makedirs(test_path, exist_ok=True) 87 | for base_path, version in zip(base_paths, versions): 88 | trials = [join(base_path, "Agent_0", e) for e in listdir(join(base_path, "Agent_0"))] 89 | # Temp 90 | random.shuffle(trials) 91 | for trial in trials: 92 | sub_dirs = listdir(join(trial)) 93 | fam_dir = [join(trial, e) for e in sub_dirs if "fam" in e.lower()] 94 | test_a = [join(trial, e) for e in sub_dirs if "_A" in e][0] 95 | test_b = [join(trial, e) for e in sub_dirs if "_B" in e][0] 96 | if not surprise_expected_reversed: 97 | surprising, expected = test_a, test_b 98 | else: 99 | surprising, expected = test_b, test_a 100 | if os.path.isdir(join(target_base_path, "test_set", f"Trial_{test_set_no}")): 101 | rmtree(join(target_base_path, "test_set", f"Trial_{test_set_no}")) 102 | makedirs(join(target_base_path, "test_set", f"Trial_{test_set_no}"), exist_ok=True) 103 | save_fam_video(fam_dir, join(target_base_path, "test_set", f"Trial_{test_set_no}")) 104 | copytree(join(trial, surprising), join(target_base_path, "test_set", f"Trial_{test_set_no}", 105 | "surprising")) 106 | copytree(join(trial, expected), join(target_base_path, "test_set", f"Trial_{test_set_no}", 107 | "expected")) 108 | open(join(target_base_path, "test_set", f"Trial_{test_set_no}", version), "w") 109 | # if os.path.isdir(join(target_base_path, "test_set", f"Set_{test_set_no}_Expected")): 110 | # rmtree(join(target_base_path, "test_set", f"Set_{test_set_no}_Expected")) 111 | # makedirs(join(target_base_path, "test_set", f"Set_{test_set_no}_Expected"), exist_ok=True) 112 | # save_fam_video(fam_dir, join(target_base_path, "test_set", f"Set_{test_set_no}_Expected")) 113 | # 114 | test_set_no += 1 115 | 116 | 117 | def process_scene_1(v1_base_path, v2_base_path, target_base_path, test_only): 118 | """ 119 | scene 1 -> scene 2 action efficiency 120 | type 1 v1 and v2 -> type 2.1 121 | type 2 v1 and v2-> type 2.2 122 | type 3 v1 -> type 2.3 123 | type 4 v2 -> type 2.4 124 | type 5 v2 -> type 2.5 125 | :return: 126 | """ 127 | 128 | makedirs(join(target_base_path, "final_dataset", "scenario_2_action_efficiency"), exist_ok=True) 129 | # Type 1 v1 and v2 130 | target_path = join(target_base_path, "final_dataset", "scenario_2_action_efficiency", "type_2.1") 131 | makedirs(target_path, exist_ok=True) 132 | create_dataset(base_paths=[join(v1_base_path, "type_1"), join(v2_base_path, "type_1")], 133 | target_base_path=target_path, test_only=test_only, size=[50, 50], versions=["version_1", "version_2"]) 134 | 135 | # Type 2 v1 and v2 136 | target_path = join(target_base_path, "final_dataset", "scenario_2_action_efficiency", "type_2.2") 137 | makedirs(target_path, exist_ok=True) 138 | create_dataset(base_paths=[join(v1_base_path, "type_2"), join(v2_base_path, "type_2")], 139 | target_base_path=target_path, test_only=test_only, size=[50, 50], versions=["version_1", "version_2"]) 140 | 141 | # Type 3 v1 142 | target_path = join(target_base_path, "final_dataset", "scenario_2_action_efficiency", "type_2.3") 143 | makedirs(target_path, exist_ok=True) 144 | create_dataset(base_paths=[join(v1_base_path, "type_3")], 145 | target_base_path=target_path, test_only=test_only, size=[100], versions=["version_1"]) 146 | 147 | # Type 4 v2 148 | target_path = join(target_base_path, "final_dataset", "scenario_2_action_efficiency", "type_2.4") 149 | makedirs(target_path, exist_ok=True) 150 | create_dataset(base_paths=[join(v2_base_path, "type_4")], 151 | target_base_path=target_path, test_only=test_only, size=[200], versions=["version_2"]) 152 | 153 | # Type 5 v2 154 | target_path = join(target_base_path, "final_dataset", "scenario_2_action_efficiency", "type_2.5") 155 | makedirs(target_path, exist_ok=True) 156 | create_dataset(base_paths=[join(v2_base_path, "type_5")], 157 | target_base_path=target_path, test_only=test_only, size=[100], versions=["version_2"]) 158 | 159 | 160 | def process_scene_2(v1_base_path, v2_base_path, target_base_path, test_only): 161 | """ 162 | scene 2 -> scene 3 unobserved constraints 163 | type 1 v1 -> type 3.1 164 | type 2 v1 & v2 -> type 3.2 165 | :return: 166 | """ 167 | 168 | makedirs(join(target_base_path, "final_dataset", "scenario_3_unobserved_constraints"), exist_ok=True) 169 | # Type 1 v1 170 | target_path = join(target_base_path, "final_dataset", "scenario_3_unobserved_constraints", "type_3.1") 171 | makedirs(target_path, exist_ok=True) 172 | create_dataset(base_paths=[join(v1_base_path, "type_1")], 173 | target_base_path=target_path, test_only=test_only, size=[200], versions=["version_1"]) 174 | 175 | # Type 2 v1 and v2 176 | target_path = join(target_base_path, "final_dataset", "scenario_3_unobserved_constraints", "type_3.2") 177 | makedirs(target_path, exist_ok=True) 178 | create_dataset(base_paths=[join(v1_base_path, "type_2"), join(v2_base_path, "type_2")], 179 | target_base_path=target_path, test_only=test_only, size=[200, 100], 180 | versions=["version_1", "version_2"]) 181 | 182 | 183 | def process_scene_3(v1_base_path, v2_base_path, target_base_path, test_only): 184 | """ 185 | scene 3 -> scene 1 Goal preferences 186 | type 1.1.0, type 1.2.0 v1 and v2 -> type 1.2 187 | type 1.0 -> type 1.1 188 | type 2.0 -> type 1.3 189 | type 2.1.0, type 2.2.0 v1 and v2 -> type 1.4 190 | :return: 191 | """ 192 | 193 | target_base_path = join(target_base_path, "final_dataset", "scenario_1_goal_preferences") 194 | makedirs(target_base_path, exist_ok=True) 195 | # Type 1.1.0 Type 1.2.0 v1 and v2 196 | target_path = join(target_base_path, "type_1.2") 197 | makedirs(target_path, exist_ok=True) 198 | create_dataset(base_paths=[join(v1_base_path, "type_1_1_0"), join(v1_base_path, "type_1_2_0"), 199 | join(v2_base_path, "type_1_2_0")], 200 | target_base_path=target_path, test_only=test_only, size=[50, 50, 100], 201 | versions=["version_1", "version_1", "version_2"], surprise_expected_reversed=True) 202 | 203 | # Type 1.0 v1 204 | target_path = join(target_base_path, "type_1.1") 205 | makedirs(target_path, exist_ok=True) 206 | create_dataset(base_paths=[join(v1_base_path, "type_1_0"), join(v2_base_path, "type_1_0")], 207 | target_base_path=target_path, test_only=test_only, size=[50, 100], 208 | versions=["version_1", "version_2"], surprise_expected_reversed=True) 209 | 210 | # Type 2.0 v1 211 | target_path = join(target_base_path, "type_1.3") 212 | makedirs(target_path, exist_ok=True) 213 | create_dataset(base_paths=[join(v1_base_path, "type_2_0"), join(v2_base_path, "type_2_0")], 214 | target_base_path=target_path, test_only=test_only, size=[50, 100], 215 | versions=["version_1", "version_2"], surprise_expected_reversed=True) 216 | 217 | # Type 2.1.0 Type 2.2.0 v1 and v2 218 | target_path = join(target_base_path, "type_1.4") 219 | makedirs(target_path, exist_ok=True) 220 | create_dataset(base_paths=[join(v1_base_path, "type_2_1_0"), join(v1_base_path, "type_2_2_0"), 221 | join(v2_base_path, "type_2_2_0")], 222 | target_base_path=target_path, test_only=test_only, size=[50, 50, 100], 223 | versions=["version_1", "version_1", "version_2"], surprise_expected_reversed=True) 224 | 225 | 226 | def process_scene_4(v1_base_path, v2_base_path, target_base_path, test_only): 227 | """ 228 | scene 4 -> scene 4 cost reward trade off 229 | type 1 v1 & v2 -> type 4.1 230 | type 2 v1 & v2 -> type 4.2 231 | :return: 232 | """ 233 | 234 | target_base_path = join(target_base_path, "final_dataset", "scenario_4_cost_reward_trade_offs") 235 | makedirs(target_base_path, exist_ok=True) 236 | # Type 1 v1 and v2 237 | target_path = join(target_base_path, "type_4.1") 238 | makedirs(target_path, exist_ok=True) 239 | create_dataset(base_paths=[join(v1_base_path, "type_1"), join(v2_base_path, "type_1")], 240 | target_base_path=target_path, test_only=test_only, size=[100, 200], 241 | versions=["version_1", "version_2"], surprise_expected_reversed=True) 242 | 243 | # Type 1 v1 and v2 244 | target_path = join(target_base_path, "type_4.2") 245 | makedirs(target_path, exist_ok=True) 246 | create_dataset(base_paths=[join(v1_base_path, "type_2"), join(v2_base_path, "type_2")], 247 | target_base_path=target_path, test_only=test_only, size=[100, 200], 248 | versions=["version_1", "version_2"], surprise_expected_reversed=True) 249 | 250 | 251 | def rename_version_files(dataset_files): 252 | scenario_list = listdir(dataset_files) 253 | for scenario in scenario_list: 254 | subtypes = listdir(join(dataset_files, scenario)) 255 | for subtype in subtypes: 256 | train_tests = listdir(join(dataset_files, scenario, subtype)) 257 | for train_test in train_tests: 258 | trials = listdir(join(dataset_files, scenario, subtype, train_test)) 259 | for trial in trials: 260 | if os.path.isfile(join(dataset_files, scenario, subtype, train_test, trial, "version_1")): 261 | os.rename(join(dataset_files, scenario, subtype, train_test, trial, "version_1"), 262 | join(dataset_files, scenario, subtype, train_test, trial, "basic")) 263 | if os.path.isfile(join(dataset_files, scenario, subtype, train_test, trial, "version_2")): 264 | os.rename(join(dataset_files, scenario, subtype, train_test, trial, "version_2"), 265 | join(dataset_files, scenario, subtype, train_test, trial, "extended")) 266 | os.rename(join(dataset_files, scenario, subtype, train_test, trial), 267 | join(dataset_files, scenario, subtype, train_test, trial.replace("Set", "Trial"))) 268 | 269 | 270 | def process_state_data(dataset_path): 271 | scenario_list = listdir(dataset_path) 272 | os.makedirs("final_dataset", exist_ok=True) 273 | for scenario in ["scenario_4_cost_reward_trade_offs"]: 274 | data = { 275 | scenario: {} 276 | } 277 | subtypes = listdir(join(dataset_path, scenario)) 278 | for subtype in subtypes: 279 | data[scenario][subtype] = {} 280 | train_tests = listdir(join(dataset_path, scenario, subtype)) 281 | for train_test in train_tests: 282 | data[scenario][subtype][train_test] = {} 283 | trials = listdir(join(dataset_path, scenario, subtype, train_test)) 284 | for trial in trials: 285 | if "Set" in trial: 286 | new_name = trial.replace("Set", "Trial") 287 | os.rename(join(dataset_path, scenario, subtype, train_test, trial), 288 | join(dataset_path, scenario, subtype, train_test, new_name)) 289 | trial = new_name 290 | data[scenario][subtype][train_test][trial] = {} 291 | sub_dirs = listdir(join(dataset_path, scenario, subtype, train_test, trial)) 292 | for sub_dir in sub_dirs: 293 | if sub_dir in ["version_1", "version_2", "basic", "extended"]: 294 | # rename the file 295 | os.rename(join(dataset_path, scenario, subtype, train_test, trial, sub_dir), 296 | join(dataset_path, scenario, subtype, train_test, trial, "version.txt")) 297 | if sub_dir == "version_1": 298 | version = "basic" 299 | elif sub_dir == "version_2": 300 | version = "extended" 301 | else: 302 | version = sub_dir 303 | sub_dir = "version.txt" 304 | # Write to the file 305 | with open(join(dataset_path, scenario, subtype, train_test, trial, sub_dir), "w") as fp: 306 | fp.write(version) 307 | data[scenario][subtype][train_test][trial]["version"] = version 308 | elif sub_dir == "version.txt": 309 | with open(join(dataset_path, scenario, subtype, train_test, trial, "version.tx"), "r") as fp: 310 | version = fp.read() 311 | data[scenario][subtype][train_test][trial]["version"] = version 312 | else: 313 | data[scenario][subtype][train_test][trial][sub_dir] = {} 314 | base_dir = join(dataset_path, scenario, subtype, train_test, trial, sub_dir) 315 | if os.path.isfile(base_dir): 316 | continue 317 | with open(join(base_dir, "state_info.json")) as fp: 318 | data[scenario][subtype][train_test][trial][sub_dir] = process_trajectory_state_data( 319 | json.load(fp), base_dir) 320 | with open(join("final_dataset", f"{scenario}.pickle"), "wb") as fp: 321 | pickle.dump(data, fp) 322 | 323 | 324 | def correct_version(dataset_path): 325 | scenario_list = listdir(dataset_path) 326 | os.makedirs("final_dataset_corrected", exist_ok=True) 327 | for scenario in scenario_list: 328 | with open(join("final_dataset", f"{scenario}.pickle"), "rb") as fp: 329 | data = pickle.load(fp) 330 | subtypes = listdir(join(dataset_path, scenario)) 331 | data_ = { 332 | scenario: {} 333 | } 334 | for subtype in subtypes: 335 | data_[scenario][subtype] = {} 336 | train_tests = listdir(join(dataset_path, scenario, subtype)) 337 | for train_test in train_tests: 338 | data_[scenario][subtype][train_test] = {} 339 | trials = listdir(join(dataset_path, scenario, subtype, train_test)) 340 | for trial in trials: 341 | if "Set" in trial: 342 | new_name = trial.replace("Set", "Trial") 343 | os.rename(join(dataset_path, scenario, subtype, train_test, trial), 344 | join(dataset_path, scenario, subtype, train_test, new_name)) 345 | trial = new_name 346 | 347 | sub_dirs = listdir(join(dataset_path, scenario, subtype, train_test, trial)) 348 | # Rename wrong keys in pickle 349 | for k in data[scenario][subtype][train_test].keys(): 350 | data_[scenario][subtype][train_test][k.replace("Set", "Trial")] = data[scenario][subtype][train_test][k] 351 | for sub_dir in sub_dirs: 352 | if sub_dir == "version.txt": 353 | version = data_[scenario][subtype][train_test][trial]["version"] 354 | with open(join(dataset_path, scenario, subtype, train_test, trial, sub_dir), "w") as fp: 355 | fp.write(version) 356 | if sub_dir == "version.tx": 357 | os.remove(join(dataset_path, scenario, subtype, train_test, trial, "version.tx")) 358 | 359 | with open(join("final_dataset_corrected", f"{scenario}.pickle"), "wb") as fp: 360 | pickle.dump(data, fp) 361 | 362 | 363 | def create_final_dataset(): 364 | if os.path.isdir("/media/data/final_dataset"): 365 | dataset_path = "/media/data/final_dataset" 366 | # makedirs(join("media", "data", "final_dataset"), exist_ok=True) 367 | # process_scene_1(v1_base_path="/media/data/AGENT/human_exp_v1/scene_1_human_exp_v1", 368 | # v2_base_path="/media/data/AGENT/human_exp_v2/scene_1_human_exp_v2", 369 | # target_base_path=os.path.join("/", "media", "data"), 370 | # test_only=True) 371 | # process_scene_2(v1_base_path="/media/data/AGENT/human_exp_v1/scene_2_human_exp_v1", 372 | # v2_base_path="/media/data/AGENT/human_exp_v2/scene_2_human_exp_v2", 373 | # target_base_path=os.path.join("/", "media", "data"), 374 | # test_only=True) 375 | # process_scene_3(v1_base_path="/media/data/AGENT/human_exp_v1/scene_3_human_exp_v1", 376 | # v2_base_path="/media/data/AGENT/human_exp_v2/scene_3_human_exp_v2", 377 | # target_base_path=os.path.join("/", "media", "data"), 378 | # test_only=True) 379 | # process_scene_4(v1_base_path="/media/data/AGENT/human_exp_v1/scene_4_human_exp_v1", 380 | # v2_base_path="/media/data/AGENT/human_exp_v2/scene_4_human_exp_v2", 381 | # target_base_path=os.path.join("/", "media", "data"), 382 | # test_only=True) 383 | else: 384 | dataset_path = "/media/data3/final_dataset" 385 | # makedirs(join("media", "data3", "final_dataset"), exist_ok=True) 386 | # process_scene_1(v1_base_path = "/media/data2/machinecommonsense/scene_1_train", 387 | # v2_base_path = "/media/data2/machinecommonsense_v2/machinecommonsense/scene_1_v2", 388 | # target_base_path=os.path.join("/", "media", "data3"), 389 | # test_only=False) 390 | # process_scene_2(v1_base_path = "/media/data2/machinecommonsense/scene_2", 391 | # v2_base_path = "/media/data2/machinecommonsense_v2/machinecommonsense/scene_2_v2", 392 | # target_base_path=os.path.join("/", "media", "data3"), 393 | # test_only=False) 394 | # process_scene_3(v1_base_path = "/media/data2/machinecommonsense/scene_3", 395 | # v2_base_path = "/media/data2/machinecommonsense_v2/machinecommonsense/scene_3_v2", 396 | # target_base_path=os.path.join("/", "media", "data3"), 397 | # test_only=False) 398 | # process_scene_4(v1_base_path = "/media/data2/machinecommonsense/scene_4", 399 | # v2_base_path = "/media/data2/machinecommonsense_v2/machinecommonsense/scene_4_v2", 400 | # target_base_path=os.path.join("/", "media", "data3"), 401 | # test_only=False) 402 | 403 | 404 | 405 | # rename_version_files(dataset_path) 406 | correct_version(dataset_path) 407 | # process_state_data(dataset_path) 408 | 409 | 410 | def validate(): 411 | # f_name = "final_dataset/scenario_1_goal_preferences.pickle" 412 | f_name = "/Users/Abhi.B@ibm.com/rotation_projects/extra_projects/cora-derenderer/proccessed_data_new/output.pickle" 413 | with open(f_name, "rb") as fp: 414 | data = pickle.load(fp) 415 | scene_key = "scenario_1_goal_preferences" 416 | for scene_sub_types in data[scene_key].keys(): 417 | for set_types in data[scene_key][scene_sub_types].keys(): 418 | for set_no in data[scene_key][scene_sub_types][set_types].keys(): 419 | sub_dirs = data[scene_key][scene_sub_types][set_types][set_no].keys() 420 | fam_data = [e for e in sub_dirs if "fam" in e.lower()] 421 | fam_data = [data[scene_key][scene_sub_types][set_types][set_no][e] for e in fam_data] 422 | expected = data[scene_key][scene_sub_types][set_types][set_no]["expected"] 423 | check_target(fam_data, expected, scene_key) 424 | 425 | 426 | def validate_versioning(): 427 | f_name = "final_dataset" 428 | scene_name = ["scenario_1_goal_preferences.pickle", "scenario_2_action_efficiency.pickle", 429 | "scenario_3_unobserved_constraints.pickle", "scenario_4_cost_reward_trade_offs.pickle"] 430 | selected_scene = scene_name[2] 431 | with open(join(f_name, selected_scene), "rb") as fp: 432 | data = pickle.load(fp) 433 | scene_key = selected_scene.replace(".pickle", "") 434 | for scene_sub_types in data[scene_key].keys(): 435 | for set_types in data[scene_key][scene_sub_types].keys(): 436 | for set_no in data[scene_key][scene_sub_types][set_types].keys(): 437 | sub_dirs = data[scene_key][scene_sub_types][set_types][set_no].keys() 438 | print(sub_dirs) 439 | 440 | 441 | def do_random_sample(): 442 | if os.path.isdir("/media/data/final_dataset"): 443 | dataset_path = "/media/data/final_dataset" 444 | else: 445 | dataset_path = "/media/data3/final_dataset" 446 | dataset_path = join(dataset_path, "scenario_4_cost_reward_trade_offs") 447 | new_path = join("random_dataset", "scenario_4_cost_reward_trade_offs") 448 | for scene_sub_types in os.listdir(dataset_path): 449 | for set_types in os.listdir(os.path.join(dataset_path, scene_sub_types)): 450 | sets = os.listdir(os.path.join(dataset_path, scene_sub_types, set_types)) 451 | random.shuffle(sets) 452 | sets = sets[0:10] 453 | for e in sets: 454 | sub_dirs = os.listdir(os.path.join(dataset_path, scene_sub_types, set_types, e)) 455 | for sub_dir in sub_dirs: 456 | if sub_dir in ["version_2", "version_1"]: 457 | continue 458 | # Copy 459 | os.makedirs(join(new_path, scene_sub_types, set_types, e, sub_dir), exist_ok=True) 460 | copyfile(join(dataset_path, scene_sub_types, set_types, e, sub_dir, "scene_c.mp4"), 461 | join(new_path, scene_sub_types, set_types, e, sub_dir, "scene_c.mp4")) 462 | 463 | shutil.make_archive(f"random_dataset.zip", 'zip', new_path) 464 | print(f"zip created at random_dataset.zip") 465 | 466 | 467 | def check_target(fam_data, expected, scene_key): 468 | if scene_key == "scenario_1_goal_preferences": 469 | assert len(fam_data) == 1 470 | fam_data = fam_data[0] 471 | agent_pos = fam_data[-1]["agent"]["center"] 472 | fam_target = 1 if agent_pos[0] > 0 else 2 473 | if fam_target == 1: 474 | fam_goal = fam_data[-1]["goal_1"]["object_label"] if fam_data[-1]["goal_1"]["center"][0] > 0 \ 475 | else fam_data[-1]["goal_2"]["object_label"] 476 | else: 477 | fam_goal = fam_data[-1]["goal_2"]["object_label"] if fam_data[-1]["goal_1"]["center"][0] > 0 \ 478 | else fam_data[-1]["goal_1"]["object_label"] 479 | agent_pos = expected[-1]["agent"]["center"] 480 | test_target = 1 if agent_pos[0] > 0 else 2 481 | if test_target == 1: 482 | expected_goal = expected[-1]["goal_1"]["object_label"] if expected[-1]["goal_1"]["center"][0] > 0 \ 483 | else expected[-1]["goal_2"]["object_label"] 484 | else: 485 | expected_goal = expected[-1]["goal_2"]["object_label"] if expected[-1]["goal_1"]["center"][0] > 0 \ 486 | else expected[-1]["goal_1"]["object_label"] 487 | assert fam_goal == expected_goal 488 | else: 489 | assert len(fam_data) == 4 490 | 491 | 492 | if __name__ == '__main__': 493 | create_final_dataset() 494 | # validate() 495 | # do_random_sample() 496 | # validate_versioning() -------------------------------------------------------------------------------- /util/barrier_with_door.py: -------------------------------------------------------------------------------- 1 | 2 | class BarrierWithDoor: 3 | def __init__(self, position, height, width, depth, door_depth, tdw_object, config_file): 4 | self.obstacle_position = position 5 | self.obstacle_height = height 6 | self.obstacle_width = width 7 | self.obstacle_depth = depth 8 | self.door_depth = door_depth 9 | self.tdw_object = tdw_object 10 | self.barrier_1, self.barrier_2, self.barrier_3 = None, None, None 11 | self.config_file = config_file 12 | self.cmd = [] 13 | 14 | def create_obstacle(self, actually_create=False,): 15 | gap = self.obstacle_depth / 2 - self.door_depth / 2 16 | self.center_1 = {"x": self.obstacle_position["x"], "y": self.obstacle_position["y"], 17 | "z": self.obstacle_position["z"] - (self.door_depth / 2 + gap / 2)} 18 | self.center_2 = {"x": self.obstacle_position["x"], "y": self.obstacle_position["y"], 19 | "z": self.obstacle_position["z"] + (self.door_depth / 2 + gap / 2)} 20 | self.center_peice_height = 0.45* self.obstacle_height if "floating" in self.config_file else 0.3*self.obstacle_height 21 | self.center_3 = {"x": self.obstacle_position["x"], "y": self.obstacle_position["y"] + self.obstacle_height/2 - 22 | self.center_peice_height/2, 23 | "z": self.obstacle_position["z"]} 24 | 25 | self.center_1_current = {"x": self.obstacle_position["x"], "y": 2.243, 26 | "z": self.obstacle_position["z"] - (self.door_depth / 2 + gap / 2)} 27 | self.center_2_current = {"x": self.obstacle_position["x"], "y": 2.243, 28 | "z": self.obstacle_position["z"] + (self.door_depth / 2 + gap / 2)} 29 | self.center_3_current = {"x": self.obstacle_position["x"], "y": 2.243 + self.obstacle_height / 2 - 30 | self.center_peice_height / 2, 31 | "z": self.obstacle_position["z"]} 32 | if "floating" in self.config_file: 33 | self.center_1_current["z"] += 10 34 | self.center_2_current["z"] += 10 35 | self.center_1["z"] += 10 36 | self.center_2["z"] += 10 37 | self.barrier_1 = self.create_barrier_( 38 | positon={"x": 10, "y": 10, "z": 10}, rotation={"x": 0, "y": 0, "z": 0}, 39 | scale={"x": self.obstacle_width, "y": self.obstacle_height, "z": gap} 40 | ) 41 | self.barrier_2 = self.create_barrier_( 42 | positon={"x": 11, "y": 10, "z": 10}, rotation={"x": 0, "y": 0, "z": 0}, 43 | scale={"x": self.obstacle_width, "y": self.obstacle_height, "z": gap} 44 | ) 45 | self.barrier_3 = self.create_barrier_( 46 | positon={"x": 12, "y": 10, "z": 10}, rotation={"x": 0, "y": 0, "z": 0}, 47 | scale={"x": self.obstacle_width, "y": self.center_peice_height, "z": self.obstacle_depth if "floating" in self.config_file else self.door_depth } 48 | ) 49 | self.gap = gap 50 | if actually_create: 51 | self.cmd = [ 52 | {"$type": "teleport_object", "position": self.center_1, "id": self.barrier_1}, 53 | {"$type": "teleport_object", "position": self.center_2, "id": self.barrier_2}, 54 | {"$type": "teleport_object", "position": self.center_3, "id": self.barrier_3} 55 | ] 56 | self.tdw_object.communicate(self.cmd) 57 | else: 58 | self.cmd = [ 59 | {"$type": "teleport_object", "position": self.center_1_current, "id": self.barrier_1}, 60 | {"$type": "teleport_object", "position": self.center_2_current, "id": self.barrier_2}, 61 | {"$type": "teleport_object", "position": self.center_3_current, "id": self.barrier_3} 62 | ] 63 | 64 | 65 | 66 | 67 | def create_barrier_(self, positon, rotation, scale): 68 | obstacle_id = self.tdw_object.add_object("prim_cube", position=positon, 69 | rotation=rotation, library="models_special.json") 70 | self.tdw_object.communicate( 71 | [{"$type": "scale_object", "scale_factor": 72 | scale, 73 | "id": obstacle_id}, 74 | 75 | {"$type": "set_mass", "id": obstacle_id, "mass": 4.0}, 76 | {"$type": "set_color", "color": 77 | {"r": 0.219607845, "g": 0.5156862754, "b": 0.2901961, "a": 1.0}, "id": obstacle_id}, 78 | # {"$type": "teleport_object", "position": positon, "id": obstacle_id}, 79 | {"$type": "set_kinematic_state", "id": obstacle_id, "is_kinematic": True, "use_gravity": False} 80 | ]) 81 | return obstacle_id 82 | 83 | def actually_create_obstacle(self): 84 | if self.cmd: 85 | self.tdw_object.communicate(self.cmd) -------------------------------------------------------------------------------- /util/camera.py: -------------------------------------------------------------------------------- 1 | class Camera: 2 | def __init__(self, camera_1_pos, camera_1_rot, camera_2_pos, camera_2_rot, camera_3_pos, camera_3_rot, scene_offset, 3 | height, width, screen_scale): 4 | self.camera_side_pos = camera_1_pos 5 | self.camera_side_rot = camera_1_rot 6 | self.camera_top_pos = camera_2_pos 7 | self.camera_top_rot = camera_2_rot 8 | self.camera_top_side_pos = camera_3_pos 9 | self.camera_top_side_rot = camera_3_rot 10 | self.current_rotation = camera_3_rot 11 | self.current_position = camera_3_pos 12 | self.camera_top_id = "a" 13 | self.camera_side_id = "b" 14 | self.camera_top_side_id = "c" 15 | self.scene_offset = scene_offset 16 | self.single_camera = False 17 | self.camera_moved = False 18 | self.height = height 19 | self.width = width 20 | self.screen_scale = screen_scale 21 | 22 | def create_camera(self, single_camera=False): 23 | self.single_camera = single_camera 24 | cmd = [ 25 | 26 | # Top camera 27 | {"$type": "create_avatar", "type": "A_Img_Caps_Kinematic", "id": self.camera_top_id}, 28 | {"$type": "teleport_avatar_to", "avatar_id": self.camera_top_id, "position": self.camera_top_pos}, 29 | {"$type": "rotate_avatar_to_euler_angles", "euler_angles": self.camera_top_rot, 30 | "avatar_id": self.camera_top_id}, 31 | 32 | {"$type": "create_avatar", "type": "A_Img_Caps_Kinematic", "id": self.camera_side_id}, 33 | {"$type": "teleport_avatar_to", "avatar_id": self.camera_side_id, "position": self.camera_side_pos}, 34 | {"$type": "rotate_avatar_to_euler_angles", "euler_angles": self.camera_side_rot, 35 | "avatar_id": self.camera_side_id}, 36 | 37 | {"$type": "create_avatar", "type": "A_Img_Caps_Kinematic", "id": self.camera_top_side_id}, 38 | {"$type": "teleport_avatar_to", "avatar_id": self.camera_top_side_id, 39 | "position": self.camera_top_side_pos}, 40 | {"$type": "rotate_avatar_to_euler_angles", "euler_angles": self.camera_top_side_rot, 41 | "avatar_id": self.camera_top_side_id}, 42 | 43 | # {"$type": "set_screen_size", "width": 1080, "height": 720}, 44 | # {"$type": "set_screen_size", "width": 480, "height": 320}, 45 | {"$type": "set_screen_size", "width": self.width, "height": self.height}, 46 | {"$type": "set_target_framerate", "framerate": 35}, 47 | # {"$type": "set_screen_size", "width": 300, "height": 200} 48 | 49 | ] 50 | if single_camera: 51 | return cmd[-5:] 52 | else: 53 | return cmd 54 | 55 | def increase_height(self, height_delta): 56 | if not self.single_camera: 57 | self.camera_side_pos["y"] += height_delta 58 | self.camera_top_pos["y"] += height_delta 59 | self.camera_top_side_pos["y"] += height_delta 60 | cmd = [ 61 | {"$type": "teleport_avatar_to", "avatar_id": self.camera_top_side_id, 62 | "position": self.camera_top_side_pos} 63 | ] 64 | if not self.single_camera: 65 | cmd.extend( 66 | [{"$type": "teleport_avatar_to", "avatar_id": self.camera_side_id, "position": self.camera_side_pos}, 67 | {"$type": "teleport_avatar_to", "avatar_id": self.camera_top_id, 68 | "position": self.camera_top_pos} 69 | ] 70 | ) 71 | return cmd 72 | 73 | def move_camera_scene_scene_3(self): 74 | aspect_ratio = [13, 8] 75 | self.width, self.height = aspect_ratio[0]*self.screen_scale, aspect_ratio[1]*self.screen_scale 76 | cmd = self.move_camera( 77 | pos={"x": 0, "y": 1.678, "z": 4.974 + self.scene_offset[2]}, 78 | rot={"x": 14.451, "y": 180, "z": 0}, 79 | avatar_id="c" 80 | ) 81 | cmd.extend([{"$type": "set_screen_size", "width": self.width, "height": self.height}]) 82 | return cmd 83 | 84 | def move_camera_scene_scene_1(self): 85 | if self.screen_scale == 40: 86 | self.width, self.height = 420, 280 87 | else: 88 | self.width, self.height = 1080, 720 89 | cmd = self.move_camera( 90 | pos={"x": 1, "y": 1.45, "z": 4.496 + self.scene_offset[2]}, 91 | rot={"x": 23.06, "y": -154.898, "z": 0}, 92 | avatar_id="c" 93 | ) 94 | cmd.extend([{"$type": "set_screen_size", "width": self.width, "height": self.height}]) 95 | return cmd 96 | 97 | def move_camera_scene_2(self): 98 | cmd = self.move_camera( 99 | pos={"x": 0.612, "y": 1.45, "z": 4.496 + self.scene_offset[2]}, 100 | rot={"x": 23.06, "y": -165.381, "z": 0}, 101 | avatar_id="c" 102 | ) 103 | self.current_position = {"x": 0.612, "y": 1.45, "z": 4.496 + self.scene_offset[2]} 104 | self.current_rotation = {"x": 23.06, "y": -165.381, "z": 0} 105 | return cmd 106 | 107 | 108 | def move_camera(self, pos, rot, avatar_id): 109 | cmd = [ 110 | {"$type": "teleport_avatar_to", "avatar_id": avatar_id, "position": pos}, 111 | {"$type": "rotate_avatar_to_euler_angles", "euler_angles": rot, 112 | "avatar_id": avatar_id}, 113 | ] 114 | self.camera_top_side_pos = pos 115 | self.camera_top_side_rot = rot 116 | self.camera_moved = True 117 | return cmd 118 | -------------------------------------------------------------------------------- /util/generate_configuration.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | RAMP_HEIGHTS = 3 4 | RAMP_ROTATIONS = 4 5 | OBSTACLE_HEIGHTS = 7 6 | OBSTACLE_WIDTH = 4 7 | OBSTACLE_DEPTH = 5 8 | AGENT_POSITION_Z = 3 9 | OBJECT_POSITON_Z = 3 10 | OBJECT_POSITON_X = 3 11 | configurations = [] 12 | # ramp_double_template = { 13 | # "barrier_type": "ramp", 14 | # "id": 0, 15 | # "agent_shape": 0, 16 | # "obj_shape": 0, 17 | # "agent_pos_z": 2, 18 | # 19 | # "ramp_height_1": 1, 20 | # "ramp_rotation_1": 3, 21 | # "obj_pos_ramp_1": 0, 22 | # "ramp_height_2": 0, 23 | # "ramp_rotation_2": 1, 24 | # "obj_pos_ramp_2": 1 25 | # } 26 | # 27 | # for ramp_height_1 in range(RAMP_HEIGHTS): 28 | # for ramp_height_2 in range(RAMP_HEIGHTS): 29 | # for ramp_rotation_1 in range(RAMP_ROTATIONS): 30 | # for ramp_rotation_2 in range(RAMP_ROTATIONS): 31 | # for agent_z in range(AGENT_POSITION_Z): 32 | # ramp_double_template["agent_pos_z"] = agent_z 33 | # ramp_double_template["ramp_rotation_1"] = ramp_rotation_1 34 | # ramp_double_template["ramp_rotation_2"] = ramp_rotation_2 35 | # ramp_double_template["ramp_height_1"] = ramp_height_1 36 | # ramp_double_template["ramp_height_2"] = ramp_height_2 37 | # configurations.append(ramp_double_template.copy()) 38 | # ramp_double_template["id"] += 1 39 | # 40 | # ramp_len = len(configurations) 41 | # print("No of ramp scenes", len(configurations)) 42 | # with open("configuration_files/configurations_3_4_ramp.json", "w") as fp: 43 | # json.dump(configurations, fp) 44 | # 45 | # configurations = [] 46 | # cube_double_template = { 47 | # "id": 0, 48 | # "barrier_type": "cube", 49 | # "agent_shape": 0, 50 | # "obj_shape": 0, 51 | # "agent_pos_z": 2, 52 | # 53 | # "obj_1_pos_x": 0, 54 | # "obj_1_pos_z": 1, 55 | # "obj_2_pos_x": 4, 56 | # "obj_2_pos_z": 1, 57 | # 58 | # "obstacle_1_pos_x": 1, 59 | # "obstacle_1_pos_z": 1, 60 | # "obstacle_1_height": 2, 61 | # "obstacle_1_width": 0, 62 | # "obstacle_1_depth": 2, 63 | # 64 | # "obstacle_2_pos_x": 3, 65 | # "obstacle_2_pos_z": 1, 66 | # "obstacle_2_height": 1, 67 | # "obstacle_2_width": 0, 68 | # "obstacle_2_depth": 2, 69 | # 70 | # "target": 0 71 | # } 72 | # 73 | # for agent_z in range(AGENT_POSITION_Z): 74 | # for obstacle_height_1 in range(6): 75 | # for obstacle_height_2 in range(6): 76 | # for obstacle_depth_1 in range(OBSTACLE_DEPTH): 77 | # for obstacle_depth_2 in range(OBSTACLE_DEPTH): 78 | # for obstacle_width_1 in range(OBSTACLE_WIDTH): 79 | # for obstacle_width_2 in range(OBSTACLE_WIDTH): 80 | # if obstacle_height_1 == 6 and obstacle_height_2 == 4: 81 | # print(cube_double_template["id"]) 82 | # cube_double_template["obstacle_1_width"] = obstacle_width_1 83 | # cube_double_template["obstacle_2_width"] = obstacle_width_2 84 | # cube_double_template["obstacle_1_depth"] = obstacle_depth_1 85 | # cube_double_template["obstacle_2_depth"] = obstacle_depth_2 86 | # cube_double_template["obstacle_1_height"] = obstacle_height_1 87 | # cube_double_template["obstacle_2_height"] = obstacle_height_2 88 | # configurations.append(cube_double_template.copy()) 89 | # cube_double_template["id"] += 1 90 | # print("No of obstacle scenes", len(configurations) ) 91 | # with open("configuration_files/configurations_3_4_cube.json", "w") as fp: 92 | # json.dump(configurations, fp) 93 | 94 | 95 | # cube_template = { 96 | # "barrier_type": "cube", 97 | # "id": 0, 98 | # "agent_shape": 0, 99 | # "obj_shape": 0, 100 | # "agent_pos_x": 0, 101 | # "agent_pos_z": 0, 102 | # "obj_pos_x": 4, 103 | # "obj_pos_z": 2, 104 | # "obstacle_pos_x": 1, 105 | # "obstacle_pos_z": 0, 106 | # "obstacle_height": 3, 107 | # "obstacle_width": 0, 108 | # "obstacle_depth": 1, 109 | # } 110 | # configurations = [] 111 | # idx = 0 112 | # for agent_x in range(3): 113 | # for agent_z in range(3): 114 | # for obstacle_pos_x in range(3): 115 | # for obstacle_pos_z in range(3): 116 | # for obstacle_height in range(OBSTACLE_HEIGHTS): 117 | # for obstacle_depth in range(OBSTACLE_DEPTH): 118 | # for obstacle_width in range(OBSTACLE_WIDTH): 119 | # for obj_pos_x in range(3): 120 | # for obj_pos_z in range(3): 121 | # if agent_x != obj_pos_x and agent_x != obstacle_pos_x and obj_pos_x != obstacle_pos_x: 122 | # cube_template["agent_pos_x"] = agent_x 123 | # cube_template["agent_pos_z"] = agent_z 124 | # cube_template["obstacle_pos_x"] = obstacle_pos_x 125 | # cube_template["obstacle_pos_z"] = obstacle_pos_z 126 | # cube_template["obstacle_height"] = obstacle_height 127 | # cube_template["obstacle_depth"] = obstacle_depth 128 | # cube_template["obstacle_width"] = obstacle_width 129 | # cube_template["obj_pos_z"] = obj_pos_z 130 | # cube_template["obj_pos_x"] = obj_pos_x 131 | # configurations.append(cube_template.copy()) 132 | # cube_template["id"] += 1 133 | # idx += 1 134 | 135 | # 136 | # ramp_template = { 137 | # "barrier_type": "ramp", 138 | # "id": 0, 139 | # "agent_shape": 1, 140 | # "obj_shape": 0, 141 | # "agent_pos_x": 0, 142 | # "agent_pos_z": 1, 143 | # "obstacle_pos_x": 1, 144 | # "obstacle_pos_z": 1, 145 | # "ramp_height": 1, 146 | # "ramp_rotation": 0, 147 | # "obj_pos_ramp": 1, 148 | # "obstacle_height": 0, 149 | # "obstacle_width": 0, 150 | # "obstacle_depth": 2, 151 | # "approach": "jump_ramp_1" 152 | # } 153 | # for agent_x in [0, 2]: 154 | # for agent_z in range(3): 155 | # for ramp_rotation in range(4): 156 | # for ramp_height in [0, 1, 3]: 157 | # for ramp_pos in [0, 1]: 158 | # if ramp_pos != 0 and agent_x != 0 or ramp_pos != 1 and agent_x != 2: 159 | # ramp_template["agent_pos_x"] = agent_x 160 | # ramp_template["agent_pos_z"] = agent_z 161 | # ramp_template["ramp_rotation"] = ramp_rotation 162 | # ramp_template["ramp_height"] = ramp_height 163 | # ramp_template["obj_pos_ramp"] = ramp_pos 164 | # configurations.append(ramp_template.copy()) 165 | # ramp_template["id"] += 1 166 | # print(len(configurations)) 167 | # with open("configuration_files/configs_single_goal_cube.json", "w") as fp: 168 | # json.dump(configurations, fp) 169 | 170 | # 171 | 172 | 173 | def create_platform_3_4_trajectories(): 174 | configurations = [] 175 | cube_template = { 176 | "barrier_type": "platform", 177 | "id": 0, 178 | "agent_shape": 0, 179 | "obj_shape": 0, 180 | "agent_pos_z": 1, 181 | "obj_pos_ramp_x_1": 0, 182 | "obj_pos_ramp_z_1": 2, 183 | "obj_pos_ramp_x_2": 0, 184 | "obj_pos_ramp_z_2": 2, 185 | "ramp_height_1": -1, 186 | "ramp_height_2": 1, 187 | } 188 | for agent_z in range(3): 189 | for ramp_height_1 in range(3): 190 | cube_template["agent_pos_z"] = agent_z 191 | cube_template["ramp_height_1"] = ramp_height_1 192 | cube_template["ramp_height_2"] = -1 193 | cube_template["id"] += 1 194 | configurations.append(cube_template.copy()) 195 | 196 | cube_template["agent_pos_z"] = agent_z 197 | cube_template["ramp_height_1"] = -1 198 | cube_template["ramp_height_2"] = ramp_height_1 199 | cube_template["id"] += 1 200 | configurations.append(cube_template.copy()) 201 | print("No of obstacle scenes", len(configurations)) 202 | with open("configuration_files/configs_platform_3_4.json", "w") as fp: 203 | json.dump(configurations, fp) 204 | 205 | 206 | def create_ramp_3_4_trajectories(): 207 | configurations = [] 208 | cube_template = { 209 | "barrier_type": "ramp", 210 | "id": 1, 211 | "agent_shape": 0, 212 | "obj_shape": 0, 213 | "agent_pos_z": 1, 214 | "ramp_height_1": 1, 215 | "ramp_rotation_1": 0, 216 | "ramp_height_2": 1, 217 | "ramp_rotation_2": 2 218 | } 219 | for agent_z in range(3): 220 | for ramp_height_1 in range(3): 221 | cube_template["agent_pos_z"] = agent_z 222 | cube_template["ramp_height_1"] = ramp_height_1 223 | cube_template["ramp_height_2"] = -1 224 | cube_template["id"] += 1 225 | configurations.append(cube_template.copy()) 226 | 227 | cube_template["agent_pos_z"] = agent_z 228 | cube_template["ramp_height_1"] = -1 229 | cube_template["ramp_height_2"] = ramp_height_1 230 | cube_template["id"] += 1 231 | configurations.append(cube_template.copy()) 232 | print("No of obstacle scenes", len(configurations)) 233 | with open("configuration_files/configs_ramp_3_4.json", "w") as fp: 234 | json.dump(configurations, fp) 235 | 236 | def create_pit_3_4_trajectories(): 237 | configurations = [] 238 | cube_template = { 239 | "barrier_type": "pit", 240 | "id": 0, 241 | "agent_shape": 0, 242 | "obj_shape": 0, 243 | "agent_pos_x": 0, 244 | "agent_pos_z": 2, 245 | "obj_1_pos_x": 0, 246 | "obj_1_pos_z": 1, 247 | "obj_2_pos_x": 4, 248 | "obj_2_pos_z": 1, 249 | "pit_width_1": 0, 250 | "pit_depth_1": 2, 251 | "pit_width_2": -1, 252 | "pit_depth_2": -1, 253 | "obstacle_1_pos_x": 1, 254 | "obstacle_1_pos_z": 1, 255 | "obstacle_2_pos_x": 3, 256 | "obstacle_2_pos_z": 1, 257 | } 258 | for agent_z in [0, 1, 2]: 259 | for obj_1_z in range(OBJECT_POSITON_Z): 260 | for pit_width_1 in range(3): 261 | cube_template["barrier_type"] = "pit" 262 | cube_template["agent_pos_z"] = agent_z 263 | cube_template["pit_width_1"] = pit_width_1 264 | cube_template["pit_width_2"] = -1 265 | cube_template["obj_1_pos_z"] = obj_1_z 266 | cube_template["id"] += 1 267 | configurations.append(cube_template.copy()) 268 | 269 | cube_template["agent_pos_z"] = agent_z 270 | cube_template["pit_width_1"] = -1 271 | cube_template["pit_width_2"] = pit_width_1 272 | cube_template["obj_2_pos_z"] = obj_1_z 273 | cube_template["id"] += 1 274 | configurations.append(cube_template.copy()) 275 | 276 | for agent_z in [0, 1, 2]: 277 | for obj_1_z in range(OBJECT_POSITON_Z): 278 | for pit_width_1 in range(3): 279 | for bridge_1_z in range(3): 280 | cube_template["barrier_type"] = "pit-with-bridge" 281 | cube_template["agent_pos_z"] = agent_z 282 | cube_template["pit_width_1"] = pit_width_1 283 | cube_template["pit_width_2"] = -1 284 | cube_template["obj_1_pos_z"] = obj_1_z 285 | cube_template["bridge_1_z"] = bridge_1_z 286 | cube_template["bridge_2_z"] = -1 287 | cube_template["id"] += 1 288 | configurations.append(cube_template.copy()) 289 | 290 | cube_template["agent_pos_z"] = agent_z 291 | cube_template["pit_width_1"] = -1 292 | cube_template["pit_width_2"] = pit_width_1 293 | cube_template["obj_2_pos_z"] = obj_1_z 294 | cube_template["bridge_1_z"] = -1 295 | cube_template["bridge_2_z"] = bridge_1_z 296 | cube_template["id"] += 1 297 | configurations.append(cube_template.copy()) 298 | 299 | print("No of obstacle scenes", len(configurations)) 300 | with open("configuration_files/configs_pit_3_4.json", "w") as fp: 301 | json.dump(configurations, fp) 302 | 303 | def create_scene_2_goal_barrier_trajectores(): 304 | configurations = [] 305 | cube_double_template = { 306 | "id": 0, 307 | "barrier_type": "cube", 308 | "agent_shape": 0, 309 | "obj_shape": 0, 310 | "agent_pos_z": 2, 311 | 312 | "obj_1_pos_x": 0, 313 | "obj_1_pos_z": 1, 314 | "obj_2_pos_x": 4, 315 | "obj_2_pos_z": 1, 316 | 317 | "obstacle_1_pos_x": 1, 318 | "obstacle_1_pos_z": 1, 319 | "obstacle_1_height": 0, 320 | "obstacle_1_width": 0, 321 | "obstacle_1_depth": 2, 322 | 323 | "obstacle_2_pos_x": 3, 324 | "obstacle_2_pos_z": 1, 325 | "obstacle_2_height": 0, 326 | "obstacle_2_width": 0, 327 | "obstacle_2_depth": 0, 328 | } 329 | 330 | for agent_z in [0, 1, 2]: 331 | for obj_1_z in range(OBJECT_POSITON_Z): 332 | for obstacle_height_1 in range(1, 6): 333 | for obstacle_width_1 in range(4): 334 | for obstacle_depth_1 in range(5): 335 | cube_double_template["agent_pos_z"] = agent_z 336 | cube_double_template["obj_1_pos_x"] = 0 337 | cube_double_template["obj_1_pos_z"] = obj_1_z 338 | cube_double_template["obstacle_1_height"] = obstacle_height_1 339 | cube_double_template["obstacle_1_width"] = obstacle_width_1 340 | cube_double_template["obstacle_1_depth"] = obstacle_depth_1 341 | cube_double_template["obstacle_2_height"] = 0 342 | configurations.append(cube_double_template.copy()) 343 | cube_double_template["id"] += 1 344 | 345 | cube_double_template["agent_pos_z"] = agent_z 346 | cube_double_template["obj_2_pos_x"] = 4 347 | cube_double_template["obj_2_pos_z"] = obj_1_z 348 | cube_double_template["obstacle_2_height"] = obstacle_height_1 349 | cube_double_template["obstacle_2_width"] = obstacle_width_1 350 | cube_double_template["obstacle_2_depth"] = obstacle_depth_1 351 | cube_double_template["obstacle_1_height"] = 0 352 | configurations.append(cube_double_template.copy()) 353 | cube_double_template["id"] += 1 354 | 355 | cube_double_template["obstacle_1_height"] = 0 356 | cube_double_template["obstacle_2_height"] = 0 357 | for agent_z in [0, 1, 2]: 358 | for obj_1_z in range(OBJECT_POSITON_Z): 359 | for obj_1_x in [0, 1, 2]: 360 | cube_double_template["agent_pos_z"] = agent_z 361 | cube_double_template["obj_1_pos_x"] = obj_1_x 362 | cube_double_template["obj_1_pos_z"] = obj_1_z 363 | configurations.append(cube_double_template.copy()) 364 | cube_double_template["id"] += 1 365 | cube_double_template["agent_pos_z"] = agent_z 366 | cube_double_template["obj_2_pos_x"] = 4 + obj_1_x 367 | cube_double_template["obj_2_pos_z"] = obj_1_z 368 | 369 | configurations.append(cube_double_template.copy()) 370 | cube_double_template["id"] += 1 371 | 372 | 373 | print("No of obstacle scenes", len(configurations) ) 374 | with open("configuration_files/configurations_3_4_cube.json", "w") as fp: 375 | json.dump(configurations, fp) 376 | 377 | def create_barrier_with_door_3_4_trajectores(): 378 | configurations = [] 379 | cube_double_template = { 380 | "id": 0, 381 | "barrier_type": "barrier_with_door", 382 | "agent_shape": 0, 383 | "obj_shape": 0, 384 | "agent_pos_z": 2, 385 | 386 | "obj_1_pos_x": 0, 387 | "obj_1_pos_z": 1, 388 | "obj_2_pos_x": 4, 389 | "obj_2_pos_z": 1, 390 | 391 | "obstacle_1_pos_x": 1, 392 | "obstacle_1_pos_z": 1, 393 | "obstacle_1_height": 6, 394 | "obstacle_1_width": 2, 395 | "obstacle_1_depth": 2, 396 | 397 | "obstacle_2_pos_x": 3, 398 | "obstacle_2_pos_z": 1, 399 | "obstacle_2_height": 6, 400 | "obstacle_2_width": 2, 401 | "obstacle_2_depth": 2, 402 | } 403 | 404 | for agent_z in range(3): 405 | for obj_1_z in range(3): 406 | for obstacle_z in range(3): 407 | cube_double_template["agent_pos_z"] = agent_z 408 | cube_double_template["obj_1_pos_x"] = 0 409 | cube_double_template["obj_1_pos_z"] = obj_1_z 410 | cube_double_template["obstacle_1_pos_z"] = obstacle_z 411 | configurations.append(cube_double_template.copy()) 412 | cube_double_template["id"] += 1 413 | 414 | cube_double_template["agent_pos_z"] = agent_z 415 | cube_double_template["obj_2_pos_x"] = 4 416 | cube_double_template["obj_2_pos_z"] = obj_1_z 417 | cube_double_template["obstacle_2_pos_z"] = obstacle_z 418 | configurations.append(cube_double_template.copy()) 419 | cube_double_template["id"] += 1 420 | 421 | 422 | print("No of obstacle scenes", len(configurations) ) 423 | with open("configuration_files/configs_barrier_with_door_3_4.json", "w") as fp: 424 | json.dump(configurations, fp) 425 | 426 | def create_ramp_trajectories(): 427 | configurations = [] 428 | cube_template = { 429 | "barrier_type": "ramp", 430 | "id": 0, 431 | "agent_shape": 0, 432 | "agent_pos_x": 2, 433 | "agent_pos_z": 1, 434 | "obj_shape": 0, 435 | "obj_pos_ramp": 0, 436 | "ramp_height": 1, 437 | "ramp_rotation": 0 438 | } 439 | for agent_x in [0, 2]: 440 | for agent_z in range(3): 441 | for obj_pos_ramp in [0, 2]: 442 | for ramp_height in range(3): 443 | for ramp_rotation in [0, 2]: 444 | if agent_x != obj_pos_ramp: 445 | if (agent_x == 0 and ramp_rotation == 2) or (agent_x == 2 and ramp_rotation == 0): 446 | cube_template["agent_pos_x"] = agent_x 447 | cube_template["agent_pos_z"] = agent_z 448 | cube_template["ramp_height"] = ramp_height 449 | cube_template["obj_pos_ramp"] = 0 if obj_pos_ramp == 0 else 1 450 | cube_template["ramp_rotation"] = ramp_rotation 451 | cube_template["id"] += 1 452 | 453 | configurations.append(cube_template.copy()) 454 | print("No of obstacle scenes", len(configurations)) 455 | with open("configuration_files/configs_ramps.json", "w") as fp: 456 | json.dump(configurations, fp) 457 | 458 | def create_platform_trajectories(): 459 | configurations = [] 460 | cube_template = { 461 | "barrier_type": "platform", 462 | "id": 0, 463 | "agent_shape": 0, 464 | "agent_pos_x": 2, 465 | "agent_pos_z": 2, 466 | "obj_shape": 0, 467 | "obj_pos_ramp_x": 0, 468 | "obj_pos_ramp_z": 2, 469 | "ramp_height": 0, 470 | } 471 | for agent_x in [0, 2]: 472 | for agent_z in range(3): 473 | for obj_pos_x in [0, 2]: 474 | for obj_pos_z in range(3): 475 | for ramp_height in range(3): 476 | if agent_x != obj_pos_x: 477 | cube_template["agent_pos_x"] = agent_x 478 | cube_template["agent_pos_z"] = agent_z 479 | cube_template["ramp_height"] = ramp_height 480 | cube_template["obj_pos_ramp_z"] = obj_pos_z 481 | cube_template["obj_pos_ramp_x"] = obj_pos_x 482 | cube_template["id"] += 1 483 | configurations.append(cube_template.copy()) 484 | print("No of obstacle scenes", len(configurations)) 485 | with open("configuration_files/configs_platform.json", "w") as fp: 486 | json.dump(configurations, fp) 487 | 488 | def create_pit_trajectories(): 489 | cube_template = { 490 | "barrier_type": "pit", 491 | "id": 0, 492 | "agent_shape": 0, 493 | "obj_shape": 0, 494 | "agent_pos_x": 0, 495 | "agent_pos_z": 0, 496 | "obj_pos_x": 4, 497 | "obj_pos_z": 2, 498 | "pit_width": 0, 499 | "pit_depth": -1, 500 | "obstacle_pos_x": 1, 501 | "obstacle_pos_z": 1, 502 | } 503 | configurations = [] 504 | idx = 0 505 | for agent_x in [0, 2]: 506 | for agent_z in range(3): 507 | for obj_pos_x in [0, 2]: 508 | for obj_pos_z in range(3): 509 | for pit_width in [-1, 0, 1, 2, 3, 4]: 510 | if agent_x != obj_pos_x: 511 | cube_template["agent_pos_x"] = agent_x 512 | cube_template["agent_pos_z"] = agent_z 513 | cube_template["pit_width"] = pit_width 514 | cube_template["obj_pos_z"] = obj_pos_z 515 | cube_template["obj_pos_x"] = obj_pos_x 516 | configurations.append(cube_template.copy()) 517 | cube_template["id"] += 1 518 | idx += 1 519 | cube_template["barrier_type"] = "pit-with-bridge" 520 | for agent_x in [0, 2]: 521 | for agent_z in range(3): 522 | for obj_pos_x in [0, 2]: 523 | for obj_pos_z in range(3): 524 | for bridge_z in range(3): 525 | for pit_width in [0, 1, 2]: 526 | if agent_x != obj_pos_x: 527 | cube_template["agent_pos_x"] = agent_x 528 | cube_template["agent_pos_z"] = agent_z 529 | cube_template["pit_width"] = pit_width 530 | cube_template["obj_pos_z"] = obj_pos_z 531 | cube_template["obj_pos_x"] = obj_pos_x 532 | cube_template["bridge_z"] = bridge_z 533 | configurations.append(cube_template.copy()) 534 | cube_template["id"] += 1 535 | idx += 1 536 | print("No of obstacle scenes", len(configurations) ) 537 | with open("configuration_files/configs_pit.json", "w") as fp: 538 | json.dump(configurations, fp) 539 | 540 | 541 | def create_barrier_trajectories(): 542 | cube_template = { 543 | "barrier_type": "cube", 544 | "id": 0, 545 | "agent_shape": 0, 546 | "obj_shape": 0, 547 | "agent_pos_x": 0, 548 | "agent_pos_z": 0, 549 | "obj_pos_x": 4, 550 | "obj_pos_z": 2, 551 | "obstacle_pos_x": 1, 552 | "obstacle_pos_z": 0, 553 | "obstacle_height": 3, 554 | "obstacle_width": 0, 555 | "obstacle_depth": 1, 556 | } 557 | configurations = [] 558 | idx = 0 559 | for agent_x in [0, 2]: 560 | for agent_z in range(3): 561 | for obstacle_pos_x in [1]: 562 | for obstacle_pos_z in range(3): 563 | for obstacle_height in range(OBSTACLE_HEIGHTS): 564 | for obstacle_depth in range(OBSTACLE_DEPTH): 565 | for obstacle_width in range(OBSTACLE_WIDTH): 566 | for obj_pos_x in [0, 2]: 567 | for obj_pos_z in range(3): 568 | if agent_x != obj_pos_x: 569 | if agent_z == obj_pos_z and obj_pos_z == obstacle_pos_z and obstacle_height != 0: 570 | print(cube_template["id"]) 571 | cube_template["agent_pos_x"] = agent_x 572 | cube_template["agent_pos_z"] = agent_z 573 | cube_template["obstacle_pos_x"] = obstacle_pos_x 574 | cube_template["obstacle_pos_z"] = obstacle_pos_z 575 | cube_template["obstacle_height"] = obstacle_height 576 | cube_template["obstacle_depth"] = obstacle_depth 577 | cube_template["obstacle_width"] = obstacle_width 578 | cube_template["obj_pos_z"] = obj_pos_z 579 | cube_template["obj_pos_x"] = obj_pos_x 580 | configurations.append(cube_template.copy()) 581 | cube_template["id"] += 1 582 | idx += 1 583 | print("No of obstacle scenes", len(configurations)) 584 | with open("configuration_files/configs_barrier.json", "w") as fp: 585 | json.dump(configurations, fp) 586 | 587 | 588 | def create_barrier_with_door_trajectories(): 589 | cube_template = { 590 | "barrier_type": "barrier_with_door", 591 | "id": 0, 592 | "agent_shape": 0, 593 | "obj_shape": 0, 594 | "agent_pos_x": 0, 595 | "agent_pos_z": 0, 596 | "obj_pos_x": 4, 597 | "obj_pos_z": 2, 598 | "obstacle_pos_x": 1, 599 | "obstacle_pos_z": 1, 600 | "obstacle_height": 6, 601 | "obstacle_width": 0, 602 | "obstacle_depth": 0, 603 | } 604 | idx = 0 605 | for agent_x in [0, 2]: 606 | for agent_z in range(3): 607 | for obj_pos_x in [0, 2]: 608 | for obj_pos_z in range(3): 609 | for obstacle_pos_z in [0, 1, 2]: 610 | for obstacle_width in [0, 1, 2, 3]: 611 | if agent_x != obj_pos_x: 612 | cube_template["agent_pos_x"] = agent_x 613 | cube_template["agent_pos_z"] = agent_z 614 | cube_template["obstacle_width"] = obstacle_width 615 | cube_template["obj_pos_z"] = obj_pos_z 616 | cube_template["obj_pos_x"] = obj_pos_x 617 | cube_template["obstacle_pos_z"] = obstacle_pos_z 618 | configurations.append(cube_template.copy()) 619 | cube_template["id"] += 1 620 | idx += 1 621 | print("No of obstacle scenes", len(configurations)) 622 | with open("configuration_files/configs_barrier_with_door.json", "w") as fp: 623 | json.dump(configurations, fp) 624 | 625 | if __name__ == '__main__': 626 | create_pit_trajectories() 627 | -------------------------------------------------------------------------------- /util/imgs/fig1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MicroSTM/AGENT-synthesis/47cc2b74a65e44ca9760a299f0eba08435130135/util/imgs/fig1.png -------------------------------------------------------------------------------- /util/imgs/icon.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MicroSTM/AGENT-synthesis/47cc2b74a65e44ca9760a299f0eba08435130135/util/imgs/icon.gif -------------------------------------------------------------------------------- /util/jump.py: -------------------------------------------------------------------------------- 1 | class jump: 2 | def __init__(self, direction=None, direction_x_z=None): 3 | self.direction = 0 if direction is None else direction 4 | self.direction_x_z = [0, 0] if direction_x_z is None else direction_x_z 5 | self.jump_forces = { 6 | "0.1_0.1": [self.direction * -2.0, 2.2, 0], 7 | "0.2_0.1": [self.direction * -1.7, 2.5, 0], 8 | "0.3_0.1": [self.direction * -1.6, 2.9, 0], 9 | "0.4_0.1": [self.direction * -1.4, 3.5, 0], 10 | "0.6_0.1": [self.direction * -1.0, 4.2, 0], 11 | "0.7_0.1": [self.direction * -1.0, 4.4, 0], 12 | "0.1_0.2": [self.direction * -2.2, 2.2, 0], 13 | "0.2_0.2": [self.direction * -1.8, 2.5, 0], 14 | "0.3_0.2": [self.direction * -1.6, 2.9, 0], 15 | "0.4_0.2": [self.direction * -1.3, 3.5, 0], 16 | "0.6_0.2": [self.direction * -1.05, 4.3, 0], 17 | "0.7_0.2": [self.direction * -1.0, 4.5, 0], 18 | "0.1_0.4": [self.direction * -2.1, 2.57, 0], 19 | "0.2_0.4": [self.direction * -1.9, 3.0, 0], 20 | "0.3_0.4": [self.direction * -1.7, 3.2, 0], 21 | "0.4_0.4": [self.direction * -1.4, 3.87, 0], 22 | "0.6_0.4": [self.direction * -1.3, 4.5, 0], 23 | "0.7_0.4": [self.direction * -1.13, 4.87, 0], 24 | "0.1_0.5": [self.direction * -2.25, 2.57, 0], 25 | "0.2_0.5": [self.direction * -2.0, 3.1, 0], 26 | "0.3_0.5": [self.direction * -1.8, 3.35, 0], 27 | "0.4_0.5": [self.direction * -1.6, 3.77, 0], 28 | "0.6_0.5": [self.direction * -1.45, 4.65, 0], 29 | "0.7_0.5": [self.direction * -1.25, 4.87, 0], 30 | "0.8_0": [direction_x_z[0] * -0.6, 4.2, direction_x_z[1] * -1.1], 31 | "0.5_0": [direction_x_z[0] * -1.1, 3.3, direction_x_z[1] * -1.1], 32 | "0.4_0": [direction_x_z[0] * -1.1, 3.0, direction_x_z[1] * -1.1], 33 | "0.2_0": [direction_x_z[0] * -1.5, 2.2, direction_x_z[1] * -1.5], 34 | "-0.7_0": [self.direction * -0.7, 4.2, 0], 35 | "-0.6_0": [self.direction * -0.7, 3.8, 0], 36 | "-0.4_0": [self.direction * -1.5, 1.1, 0], 37 | "-0.3_0": [self.direction * -0.9, 2.7, 0], 38 | "-0.2_0": [self.direction * -1.5, 1.1, 0], 39 | "-0.1_0": [self.direction * -0.9, 2.2, 0], 40 | "None_0.38": [self.direction * -1.7, 2, 0], 41 | "None_0.6": [self.direction * -2.2, 2.1, 0], 42 | "None_0.88": [self.direction * -2.7, 2.2, 0], 43 | "None_1.13": [self.direction * -2.9, 2.5, 0], 44 | "None_1.38": [self.direction * -3.1, 2.7, 0], 45 | "None_0.7": [self.direction * -1.9, 2.5, 0], 46 | "None_0.5": [self.direction * -1.7, 2.3, 0], 47 | "None_0.3": [self.direction * -1.4, 2.2, 0], 48 | } 49 | 50 | def get_jump_force(self, agent_velocity, jump_height, jump_width): 51 | jump_key = f"{jump_height}_{jump_width}" 52 | jump_force = self.jump_forces[jump_key] 53 | return [e - f for e, f in zip(jump_force, agent_velocity)] -------------------------------------------------------------------------------- /util/obstacle.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | 4 | class Obstacle: 5 | def __init__(self, tdw_object, position, height, width, rotation={"x": 0, "y": 0, "z": 0}, depth=2.7, door_depth=None, config_file=None): 6 | self.tdw_object = tdw_object 7 | self.obstacle_height = height 8 | self.obstacle_width = width 9 | self.obstacle_depth = depth 10 | self.obstacle_position = position 11 | self.obstacle_id = None 12 | self.current_position = None 13 | self.obstacle_rotation = rotation 14 | self.colors = [{"r": 0.219607845, "g": 0.5156862754, "b": 0.2901961, "a": 1.0}] 15 | self.cmd = [] 16 | 17 | def create_obstacle(self, actually_create=False): 18 | 19 | 20 | self.obstacle_id = self.tdw_object.add_object("prim_cube", position={"x": 10, "y": 10, "z": 10}, 21 | rotation=self.obstacle_rotation, library="models_special.json") 22 | self.current_position = {"x": self.obstacle_position["x"], "y": 2.243, "z": self.obstacle_position["z"]} 23 | self.tdw_object.communicate([ 24 | {"$type": "set_kinematic_state", "id": self.obstacle_id, "is_kinematic": True, 25 | "use_gravity": False}, 26 | {"$type": "scale_object", "scale_factor": 27 | {"x": self.obstacle_width, "y": self.obstacle_height, "z": self.obstacle_depth}, 28 | "id": self.obstacle_id}, 29 | {"$type": "set_mass", "id": self.obstacle_id, "mass": 4.0}, 30 | {"$type": "set_color", "color": 31 | {"r": 0.219607845, "g": 0.5156862754, "b": 0.2901961, "a": 1.0}, "id": self.obstacle_id}, 32 | ]) 33 | 34 | if actually_create: 35 | self.tdw_object.communicate({"$type": "teleport_object", "position": self.obstacle_position, "id": self.obstacle_id}) 36 | else: 37 | self.cmd = [ 38 | {"$type": "teleport_object", "position": self.current_position, "id": self.obstacle_id}, 39 | 40 | ] 41 | 42 | return self.obstacle_id 43 | 44 | def actually_create_obstacle(self): 45 | if self.cmd: 46 | self.tdw_object.communicate(self.cmd) 47 | 48 | -------------------------------------------------------------------------------- /util/occuluder.py: -------------------------------------------------------------------------------- 1 | from util import utils 2 | from util.utils import SceneConfiguration 3 | import numpy as np 4 | import util.trajectory_generation 5 | 6 | class Occluder: 7 | def __init__(self, camera_obj, obstacle, pit=False): 8 | self.scene_config = SceneConfiguration() 9 | self.camera_obj = camera_obj 10 | self.obstacle = obstacle 11 | self.reveal_normal = [] 12 | self.calculate_occuluder_properties(pit) 13 | 14 | def y_rotation(self, vector, theta): 15 | """Rotates 3-D vector around y-axis""" 16 | R = np.array([[np.cos(theta), 0, np.sin(theta)], [0, 1, 0], [-np.sin(theta), 0, np.cos(theta)]]) 17 | return np.dot(R, vector) 18 | 19 | def calculate_occuluder_properties(self, pit=False): 20 | camera_pos = self.camera_obj.current_position.copy() 21 | camera_rot = self.camera_obj.current_rotation.copy() 22 | 23 | # Find center of obstacle 24 | front_right_point = np.array([ 25 | self.obstacle.obstacle_position["x"] - self.obstacle.obstacle_width/2, 26 | self.obstacle.obstacle_position["y"] + self.obstacle.obstacle_height/2, 27 | self.obstacle.obstacle_position["z"] + self.obstacle.obstacle_depth / 2 28 | ]) 29 | back_left_point = np.array([ 30 | self.obstacle.obstacle_position["x"] + self.obstacle.obstacle_width / 2, 31 | self.obstacle.obstacle_position["y"] - self.obstacle.obstacle_height / 2, 32 | self.obstacle.obstacle_position["z"] - self.obstacle.obstacle_depth / 2 33 | ]) 34 | back_right_point = np.array([ 35 | self.obstacle.obstacle_position["x"] - self.obstacle.obstacle_width / 2, 36 | self.obstacle.obstacle_position["y"] + self.obstacle.obstacle_height / 2, 37 | self.obstacle.obstacle_position["z"] - self.obstacle.obstacle_depth / 2 38 | ]) 39 | 40 | normal_vect = self.y_rotation(np.array([0, 0, 1]), np.radians(camera_rot["y"])) 41 | normal_vect_normal = self.y_rotation(normal_vect, np.radians(-90)) 42 | self.reveal_normal = -1*normal_vect 43 | projected_point = np.array([ 44 | [front_right_point[0] + 5 * normal_vect_normal[0], front_right_point[2] + 5 * normal_vect_normal[2]], 45 | [front_right_point[0], front_right_point[2]] 46 | ]) 47 | p = trajectory_generation.line_intersection( 48 | projected_point, 49 | [ 50 | [camera_pos["x"], camera_pos["z"]], 51 | [back_left_point[0], back_left_point[2]], 52 | ] 53 | ) 54 | vect_to_target = np.array([ 55 | (p[0] + front_right_point[0]) / 2 - camera_pos["x"], 56 | (p[1] + front_right_point[2]) / 2 - camera_pos["z"] 57 | ]) 58 | vect_to_target = vect_to_target/np.linalg.norm(vect_to_target) 59 | self.occluder_position = { 60 | "x": camera_pos["x"] + vect_to_target[0], 61 | "y": 0, 62 | "z": camera_pos["z"] + vect_to_target[1] 63 | } 64 | dist = utils.distance_of_point_from_line(np.array([camera_pos["x"], camera_pos["z"]]), 65 | np.array([front_right_point[0], front_right_point[2]]), 66 | np.array([self.occluder_position["x"], self.occluder_position["z"]])) 67 | 68 | height = trajectory_generation.line_intersection( 69 | [ 70 | [camera_pos["y"], camera_pos["z"]], 71 | [back_right_point[1], back_right_point[2]] 72 | ], 73 | [ 74 | [self.occluder_position["y"], self.occluder_position["z"]], 75 | [self.occluder_position["y"] + 5*1, self.occluder_position["z"]], 76 | ] 77 | )[0] 78 | x_correction = 0.023 79 | self.occluder_position["y"] = height/2 80 | self.occluder_position["y"] += 0 if not pit else 1.009 81 | self.scale = {"x": 2*dist + x_correction, "y": height, "z": 0.01} 82 | self.occluder_rotation = {"x":0, "y":camera_rot["y"], "z":0} 83 | print(self.occluder_position) 84 | 85 | 86 | def create(self, tdw_object): 87 | self.occluder_id = utils.create_occluder(tdw_object, self.occluder_position, self.scale, self.occluder_rotation) 88 | 89 | # tdw_object.communicate( 90 | # {"$type": "set_kinematic_state", "id": self.occluder_id, "is_kinematic": True, 91 | # "use_gravity": False} 92 | # ) -------------------------------------------------------------------------------- /util/onegoalscene.py: -------------------------------------------------------------------------------- 1 | from util import physics_agent, utils 2 | import random 3 | import shutil 4 | import os 5 | from util.utils import SceneConfiguration 6 | import json 7 | import numpy as np 8 | import time 9 | import traceback 10 | from util import trajectory_generation 11 | from util import create_scene 12 | 13 | 14 | class OneGoalScene: 15 | def __init__(self, tdw_object, camera_obj, args): 16 | self.tdw_object = tdw_object 17 | self.args = args 18 | self.direction = random.sample([-1, 1], 1)[0] 19 | self.scene_config = SceneConfiguration() 20 | self.scene_offset = self.scene_config.scene_offset 21 | self.obstacle_height = random.sample(self.scene_config.obstacle_height, 1)[0] 22 | self.obstacle_width = random.sample(self.scene_config.obstacle_width, 1)[0] 23 | 24 | self.agent_type = "cone" 25 | self.agent_color = random.sample(self.scene_config.agent_colors, 1)[0] 26 | self.target_color = random.sample(self.scene_config.target_colors, 1)[0] 27 | self.agent_color = {"r": 7 / 255, "g": 121 / 255, "b": 228 / 255, "a": 1.0} 28 | self.target_color = {"r": 255 / 255, "g": 87 / 255, "b": 51 / 255, "a": 1.0} 29 | # self.target_color = {"r": 133 / 255, "g": 102 / 255, "b": 170 / 255, "a": 1.0} 30 | self.config_file = None 31 | self.obstacle = None 32 | self.pit_obj = None 33 | self.camera_obj = camera_obj 34 | self.high_scene = False 35 | self.ramp = None 36 | self.pit_ids = [] 37 | self.results = [] 38 | self.idx = 0 39 | 40 | def execute_path(self, path): 41 | idx = 0 42 | meta_data = None 43 | for i, point in enumerate(path): 44 | if point[0] == "goto": 45 | if len(point) == 6: 46 | angle = self.agent.calculate_angle([point[1], point[3]]) 47 | 48 | velocity_threshold = point[5] 49 | self.agent.approach(till_x=point[1], 50 | till_y=point[2], 51 | till_z=point[3], approach_dist=point[4], speed_threshold=velocity_threshold, 52 | rotation=angle) 53 | elif len(point) == 5: 54 | angle = self.agent.calculate_angle([point[1], point[2]]) 55 | 56 | velocity_threshold = point[4] 57 | self.agent.approach(till_x=point[1], 58 | till_z=point[2], approach_dist=point[3], speed_threshold=velocity_threshold, 59 | rotation=angle) 60 | else: 61 | if self.config_file["barrier_type"] == "pit": 62 | angle = self.agent.calculate_angle([path[-2][1], path[-2][2]]) 63 | else: 64 | angle = self.agent.calculate_angle([point[1], point[2]]) 65 | velocity_threshold = 1.3 66 | self.agent.approach(till_x=point[1], 67 | till_z=point[2], approach_dist=point[3], speed_threshold=velocity_threshold, 68 | rotation=angle) 69 | elif point[0] == "jump": 70 | 71 | # look for next goto point and turn towards that 72 | for e in path[i:]: 73 | if e[0] == "goto": 74 | if len(e) == 6: 75 | angle = self.agent.calculate_angle([e[1], e[3]]) 76 | else: 77 | angle = self.agent.calculate_angle([e[1], e[2]]) 78 | self.agent.rotate_by_angle(angle) 79 | break 80 | if self.config_file["barrier_type"] in ["ramp", "platform", "platform-cube"]: 81 | direction_x_z = np.array([-self.target_position["x"] + self.agent.agent_position["x"], 82 | -self.target_position["z"] + self.agent.agent_position["z"]]) 83 | direction_x_z = direction_x_z / np.linalg.norm(direction_x_z) 84 | self.agent.jump(height=point[1], width=point[2], direction_x_z=direction_x_z, hight_ground_height=point[1]) 85 | else: 86 | self.agent.jump(height=point[1], width=point[2]) 87 | elif point[0] == "settle": 88 | self.agent.settle_drop_object(point[1], stop_at_base=False) 89 | elif point[0] == "meta": 90 | meta_data = point 91 | idx += 1 92 | 93 | # self.agent.reveal_occluder(self.occluder) 94 | 95 | # utils.make_video(self.agent.output_dir + "/images" + "_a", f"scene_a.mp4", os.path.join(self.agent.base_dir, f"path_{self.agent.path_no}")) 96 | # utils.make_video(self.agent.output_dir + "/images" + "_b", f"scene_b.mp4", os.path.join(self.agent.base_dir, f"path_{self.agent.path_no}")) 97 | if self.args.enable_image: 98 | utils.make_video(self.agent.output_dir + "_c", f"scene_c.mp4", os.path.join(self.agent.base_dir, f"path_{self.agent.path_no}")) 99 | # shutil.rmtree(self.agent.output_dir + "/images" + "_a") 100 | # shutil.rmtree(self.agent.output_dir + "/images" + "_b") 101 | shutil.rmtree(self.agent.output_dir + "_c") 102 | # Save data 103 | 104 | # Create object meta 105 | object_meta = create_scene.construct_1_goal_object_meta_data( 106 | obstacle=self.obstacle, target_id=self.target_id, 107 | target_color=self.target_color, 108 | agent_color=self.agent_color, 109 | target_position=self.target_position, 110 | ramp=self.ramp, config_file=self.config_file, 111 | agent=self.agent, pit_ids=self.pit_ids, 112 | occluder=None, 113 | pit_obj=self.pit_obj 114 | ) 115 | state_data = create_scene.create_state_data(self.agent, object_meta, self.camera_obj, meta_data) 116 | 117 | self.total_force = [0, 0, 0] 118 | for i, e in enumerate(self.agent.forces): 119 | self.total_force[0] += abs(e["x"]) 120 | self.total_force[1] += abs(e["y"]) 121 | self.total_force[2] += abs(e["z"]) 122 | 123 | self.results.append((self.total_force, sum(self.total_force), meta_data)) 124 | file_path = os.path.join(self.agent.base_dir, f"path_{self.agent.path_no}") 125 | os.makedirs(file_path, exist_ok=True) 126 | with open(os.path.join(file_path, "state_info.json"), "w") as fp: 127 | json.dump(state_data, fp) 128 | 129 | def execute_config(self, config_file, output_dir): 130 | self.config_file = config_file 131 | self.agent = physics_agent.Agent(self.tdw_object, 1, config_file["dir_name"]) 132 | self.target_shape = "sphere" 133 | 134 | # self.target_color = self.scene_config.target_colors[self.idx] 135 | # self.idx += 1 136 | return_state = create_scene.create_based_config(config_file=config_file, tdw_object=self.tdw_object, 137 | high_scene=self.high_scene, camera_obj=self.camera_obj, 138 | agent_color=self.agent_color, scene_config=self.scene_config, 139 | agent_class=self.agent, target_color=self.target_color, 140 | agent_shape=self.agent_type, 141 | enable_image=self.args.enable_image, target_shape=self.target_shape, 142 | actually_create=True, args=self.args) 143 | # add_occuluder=True, scene_state_old=self.config_file) 144 | self.agent = return_state["agent"] 145 | 146 | self.obstacle = return_state["obstacle"] 147 | self.pit_ids = return_state["pit_ids"] 148 | self.ramp = return_state["ramp"] 149 | self.high_scene = return_state["high_scene"] 150 | self.target_position = return_state["target_position"] 151 | self.target_id = return_state["target_id"] 152 | self.direction = return_state["direction"] 153 | self.pit_obj = return_state["pit_obj"] 154 | # self.occluder = return_state["occluder"] 155 | try: 156 | paths = self.go_to_goal() 157 | # paths = [ 158 | # [("rotate", 90), ("rotate", -90)] 159 | # ] 160 | 161 | results = [] 162 | time_required = [] 163 | # materials = utils.select_material(self.args) 164 | # utils.set_wall_floor_material(self.tdw_object, materials["floor_material"], 165 | # materials["wall_material"], pit_obj=self.pit_obj 166 | # ) 167 | for i, p in enumerate(paths): 168 | start_time = time.time() 169 | utils.settle_objects(self.tdw_object, n=30) 170 | self.agent.output_dir = os.path.join(self.agent.output_dir, "images") 171 | # self.agent.look_at_goal() 172 | resp = self.tdw_object.communicate({"$type": "do_nothing"}) 173 | self.agent.forces.append({"x":0, "y":0, "z":0}) 174 | self.agent.save_imag(resp) 175 | self.execute_path(p) 176 | time_required.append(time.time() - start_time) 177 | print(f"Took {time.time() - start_time} secs. Average time per trajectory {sum(time_required)/len(time_required)}") 178 | unique_str = config_file["dir_name"] + f"_path_no_{self.agent.path_no}" 179 | with open(os.path.join(output_dir, "time.txt"), "a") as fp: 180 | fp.write(f"{unique_str} took {time.time() - start_time} secs. Average time per trajectory {sum(time_required)/len(time_required)} \n") 181 | if self.config_file["barrier_type"] == "ramp": 182 | results.append( ( 183 | self.ramp.rotation, f"Agent path number {self.agent.path_no}", self.agent.total_force.copy() 184 | )) 185 | elif self.config_file["barrier_type"] == "cube": 186 | results.append((f"Agent path number {self.agent.path_no}", self.agent.total_force)) 187 | self.reset() 188 | except Exception as err: 189 | with open("error.txt", "a") as fp: 190 | fp.write(f"Error in configuration {output_dir}\n") 191 | traceback.print_tb(err.__traceback__) 192 | print(str(err)) 193 | with open("out.txt", "a") as fp: 194 | fp.write(f"-----------------\n") 195 | for e in self.results: 196 | with open("out.txt", "a") as fp: 197 | fp.write(f"{e}\n") 198 | with open("out.txt", "a") as fp: 199 | fp.write(f"-----------------\n") 200 | self.results = [] 201 | os.makedirs(self.config_file["dir_name"], exist_ok=True) 202 | with open(os.path.join(self.config_file["dir_name"], "scene_config.json"), "w") as fp: 203 | json.dump(self.config_file, fp) 204 | # Destroy scene objects 205 | self.destroy_objects() 206 | self.ramp = None 207 | self.obstacle = None 208 | self.agent = None 209 | # for r in results: 210 | # print(r) 211 | 212 | 213 | def destroy_objects(self): 214 | commands = [{"$type": "destroy_object", "id": self.agent.agent_id}, 215 | {"$type": "destroy_object", "id": self.target_id} 216 | ] 217 | if self.config_file["barrier_type"] == "barrier_with_door": 218 | commands.append({"$type": "destroy_object", "id": self.obstacle.barrier_1}) 219 | commands.append({"$type": "destroy_object", "id": self.obstacle.barrier_2}) 220 | commands.append({"$type": "destroy_object", "id": self.obstacle.barrier_3}) 221 | elif self.config_file["barrier_type"] == "cube": 222 | if self.obstacle.obstacle_id is not None: 223 | commands.append({"$type": "destroy_object", "id": self.obstacle.obstacle_id}) 224 | elif self.config_file["barrier_type"] == "ramp": 225 | commands.extend([ 226 | {"$type": "destroy_object", "id": self.ramp.ramp_base_id}, 227 | {"$type": "destroy_object", "id": self.ramp.ramp_slope_id} 228 | ]) 229 | elif self.config_file["barrier_type"] == "platform": 230 | commands.extend([ 231 | {"$type": "destroy_object", "id": self.ramp.ramp_base_id} 232 | ]) 233 | for pit_id in self.pit_ids: 234 | commands.append( 235 | {"$type": "destroy_object", "id": pit_id} 236 | ) 237 | self.pit_ids = [] 238 | resp = self.tdw_object.communicate(commands) 239 | utils.print_log(resp) 240 | 241 | def go_to_goal(self): 242 | if self.config_file["barrier_type"] in ["pit", "pit-with-bridge"]: 243 | return trajectory_generation.cross_pit(self.agent, self.config_file, self.target_position, self.pit_obj) 244 | elif self.config_file["barrier_type"] == "cube": 245 | # Find a way around or over barrier 246 | return trajectory_generation.cross_barrier(self.obstacle, self.target_position, self.agent.agent_start_position, 247 | self.scene_config, config_file=self.config_file) 248 | elif self.config_file["barrier_type"] == "barrier_with_door": 249 | return trajectory_generation.cross_barrier_with_door(self.obstacle, self.target_position, 250 | self.agent.agent_start_position) 251 | elif self.config_file["barrier_type"] == "platform-cube": 252 | return trajectory_generation.cross_platform_cube(self.obstacle, self.target_position, 253 | self.agent.agent_start_position) 254 | elif self.config_file["barrier_type"] == "ramp": 255 | paths = trajectory_generation.cross_ramp(self.ramp, self.agent, self.config_file, 256 | restrict_paths=[1, 2, 3]) 257 | return paths 258 | elif self.config_file["barrier_type"] == "platform": 259 | paths = trajectory_generation.cross_platform(self.ramp, self.agent) 260 | return paths 261 | else: 262 | return [ 263 | [("goto", self.target_position["x"],self.target_position["z"], 0.24)] 264 | ] 265 | 266 | def reset(self): 267 | self.agent.reset() 268 | utils.stop_obj(self.tdw_object, self.target_id) 269 | utils.teleport(self.tdw_object, self.target_position, self.target_id) 270 | -------------------------------------------------------------------------------- /util/physics_agent.py: -------------------------------------------------------------------------------- 1 | from util.utils import euler_to_quaternion, convert_to_radians 2 | import os 3 | from tdw.tdw_utils import TDWUtils 4 | from tdw.output_data import Images, OutputData, Transforms, Rigidbodies, Bounds, SegmentationColors 5 | import math 6 | import numpy as np 7 | from util import utils 8 | from util.jump import jump as jump_forces_list 9 | 10 | 11 | class Agent: 12 | def __init__(self, tdw_object, direction, out_dir): 13 | self.tdw_object = tdw_object 14 | self.direction = direction 15 | self.idx = 0 16 | self.path_no = 1 17 | self.output_dir = os.path.join(out_dir, f"path_{self.path_no}") 18 | self.base_dir = out_dir 19 | self.agent_positions = [] 20 | self.agent_position = None 21 | self.agent_rotations = [] 22 | self.agent_velocities = [] 23 | self.total_force = [0, 0, 0] 24 | self.forces = [] 25 | self.state_data = [] 26 | self.agent_start_position = None 27 | self.agent_angle = 0 28 | self.view_vector = [0, 1] 29 | self.image_data = {} 30 | self.agent_state = { 31 | "velocity": [0, 0, 0] 32 | } 33 | self.agent_state_data = [] 34 | 35 | def reset(self): 36 | self.stop_agent() 37 | utils.teleport(self.tdw_object, self.agent_start_position, self.agent_id) 38 | utils.rotate_object(self.tdw_object, {"x": 0, "y": 0, "z": 0}, self.agent_id) 39 | self.agent_positions = [] 40 | self.agent_position = self.agent_start_position.copy() 41 | self.agent_rotations = [] 42 | self.agent_velocities = [] 43 | self.state_data = [] 44 | self.agent_rotation = {"x": 0, "y": 0, "z": 0} 45 | self.total_force = [0, 0, 0] 46 | self.forces = [] 47 | self.view_vector = [0, 1] 48 | self.agent_angle = 0 49 | self.idx = 0 50 | self.path_no += 1 51 | self.image_data = {} 52 | self.agent_state_data = [] 53 | self.output_dir = os.path.join(self.base_dir, f"path_{self.path_no}") 54 | 55 | def calculate_angle(self, position): 56 | vect_to_target = [position[0] - self.agent_position["x"], 57 | position[1] - self.agent_position["z"]] 58 | vect_to_target = vect_to_target/np.linalg.norm(vect_to_target) 59 | angle = np.degrees(np.math.atan2(np.linalg.det([self.view_vector, vect_to_target]), np.dot(self.view_vector, vect_to_target))) 60 | self.view_vector = vect_to_target 61 | angle = -1*int(round(angle)) 62 | return angle 63 | 64 | def create_agent_new(self, position, color, agent_name): 65 | selected_color = color 66 | # agent_name = "cone" 67 | # agent_name = available_agents[random.randint(0, len(available_agents) - 1)] 68 | url = "file:///" + os.path.join(os.getcwd(), "../agent_models", agent_name.title(), "StandaloneOSX", 69 | agent_name) 70 | # url = f"file:///agent_model/{agent_name.title()}/StandaloneLinux64/{agent_name}" 71 | self.agent_id = self.tdw_object.get_unique_id() 72 | self.agent_start_position = position.copy() 73 | self.agent_position = position 74 | self.agent_rotation = {"x": 0, "y": 0, "z": 0} 75 | # self.tdw_object.communicate({"$type": "add_object", 76 | # "name": "cone", 77 | # "url": url, 78 | # "scale_factor": 1, 79 | # "id": self.agent_id, 80 | # "position": position}) 81 | self.agent_id = self.tdw_object.add_object("prim_cone", position=self.agent_position, rotation=self.agent_rotation, 82 | library="models_special.json") 83 | resp = self.tdw_object.communicate([ 84 | {"$type": "scale_object", "scale_factor": {"x": 0.2, "y": 0.2, "z": 0.2}, "id": self.agent_id}, 85 | {"$type": "set_color_in_substructure", 86 | "color": selected_color, 87 | "object_name": "prim_cone", "id": self.agent_id}, 88 | {"$type": "set_mass", "id": self.agent_id, "mass": 1.0}, 89 | {"$type": "step_physics", "frames": 10}, 90 | {"$type": "set_physic_material", "dynamic_friction": 0.0, "static_friction": 0.0, "bounciness": 0.125, 91 | "id": self.agent_id}, 92 | {"$type": "set_object_drag", "id": self.agent_id, "drag": 0, "angular_drag": 100} 93 | ]) 94 | return 95 | 96 | def look_at_goal(self): 97 | for i in range(0, self.direction * -90, self.direction * -2): 98 | qo = euler_to_quaternion(0, convert_to_radians(i), 0) 99 | resp = self.tdw_object.communicate( 100 | {"$type": "rotate_object_to", "rotation": {"w": qo[3], "x": qo[0], "y": qo[1], "z": qo[2]}, 101 | "id": self.agent_id}) 102 | 103 | self.forces.append({"x": 0.0, "y": 0.0, "z": 0.0}) 104 | self.save_imag(resp) 105 | self.agent_angle = i 106 | 107 | def rotate_by_angle(self, angle): 108 | ang = 0 109 | return 110 | direction = 1 if angle < 0 else -1 111 | for i in range(self.agent_angle, self.agent_angle + angle, direction * -5): 112 | qo = euler_to_quaternion(0, convert_to_radians(i), 0) 113 | resp = self.tdw_object.communicate( 114 | {"$type": "rotate_object_to", "rotation": {"w": qo[3], "x": qo[0], "y": qo[1], "z": qo[2]}, 115 | "id": self.agent_id}) 116 | 117 | self.forces.append({"x": 0.0, "y": 0.0, "z": 0.0}) 118 | self.save_imag(resp) 119 | ang = i 120 | self.agent_angle = ang 121 | 122 | def look_back(self): 123 | for i in range(self.direction * -90, 0, self.direction * 2): 124 | self.agent_rotation = {"x": 0, "y": i, "z": 0} 125 | 126 | qo = euler_to_quaternion(0, convert_to_radians(i), 0) 127 | resp = self.tdw_object.communicate( 128 | {"$type": "rotate_object_to", "rotation": {"w": qo[3], "x": qo[0], "y": qo[1], "z": qo[2]}, 129 | "id": self.agent_id}) 130 | self.forces.append({"x": 0.0, "y": 0.0, "z": 0.0}) 131 | self.save_imag(resp) 132 | 133 | def get_agent_loc_vel(self, response, get_x=True, both=False, send_rotation=False): 134 | velocity, position = None, None 135 | for r in response: 136 | r_id = OutputData.get_data_type_id(r) 137 | if r_id == "tran": 138 | transform_data = Transforms(r) 139 | for object_index in range(transform_data.get_num()): 140 | if transform_data.get_id(object_index) == self.agent_id: 141 | pos = transform_data.get_position(object_index) 142 | position = {"x": pos[0], "y": pos[1], "z": pos[2]} 143 | rotation = transform_data.get_rotation(object_index) 144 | if r_id == "rigi": 145 | rigid_body_data = Rigidbodies(r) 146 | for object_index in range(rigid_body_data.get_num()): 147 | if rigid_body_data.get_id(object_index) == self.agent_id: 148 | if both: 149 | velocity = (rigid_body_data.get_velocity(object_index)[0], 150 | rigid_body_data.get_velocity(object_index)[1], 151 | rigid_body_data.get_velocity(object_index)[2]) 152 | elif get_x: 153 | velocity = rigid_body_data.get_velocity(object_index)[0] 154 | else: 155 | velocity = rigid_body_data.get_velocity(object_index)[2] 156 | 157 | if send_rotation: 158 | return velocity, position, rotation 159 | else: 160 | return velocity, position 161 | 162 | def save_imag(self, response): 163 | state_data = {} 164 | for r in response: 165 | r_id = OutputData.get_data_type_id(r) 166 | if r_id == "imag": 167 | images = Images(r) 168 | image_dir = self.output_dir + "_" + images.get_avatar_id() 169 | TDWUtils.save_images(images, f"image-{self.idx}", 170 | output_directory=image_dir) 171 | if r_id == "segm": 172 | seg_data = SegmentationColors(r) 173 | for i in range(seg_data.get_num()): 174 | if seg_data.get_object_id(i) not in state_data: 175 | state_data[seg_data.get_object_id(i)] = {} 176 | state_data[seg_data.get_object_id(i)]["segmentation_color"] = seg_data.get_object_color(i) 177 | if r_id == "tran": 178 | transform_data = Transforms(r) 179 | for object_index in range(transform_data.get_num()): 180 | if transform_data.get_id(object_index) not in state_data: 181 | state_data[transform_data.get_id(object_index)] = {} 182 | state_data[transform_data.get_id(object_index)].update({ 183 | "position": transform_data.get_position(object_index), 184 | "rotation": transform_data.get_rotation(object_index) 185 | }) 186 | if transform_data.get_id(object_index) == self.agent_id: 187 | pos = transform_data.get_position(object_index) 188 | self.agent_position = {"x": pos[0], "y": pos[1], "z": pos[2]} 189 | self.agent_rotation = transform_data.get_rotation(object_index) 190 | if r_id == "rigi": 191 | rigid_body_data = Rigidbodies(r) 192 | for object_index in range(rigid_body_data.get_num()): 193 | if rigid_body_data.get_id(object_index) == self.agent_id: 194 | self.agent_state["velocity"] = (rigid_body_data.get_velocity(object_index)[0], 195 | rigid_body_data.get_velocity(object_index)[1], 196 | rigid_body_data.get_velocity(object_index)[2]) 197 | if rigid_body_data.get_id(object_index) not in state_data: 198 | state_data[rigid_body_data.get_id(object_index)] = {} 199 | state_data[rigid_body_data.get_id(object_index)].update({ 200 | "velocity": rigid_body_data.get_velocity(object_index), 201 | "angular_velocity": rigid_body_data.get_angular_velocity(object_index) 202 | }) 203 | 204 | if r_id == "boun": 205 | bounding_box_data = Bounds(r) 206 | for object_index in range(bounding_box_data.get_num()): 207 | if bounding_box_data.get_id(object_index) not in state_data: 208 | state_data[bounding_box_data.get_id(object_index)] = {} 209 | state_data[bounding_box_data.get_id(object_index)]["bounding_box"] = { 210 | "center": bounding_box_data.get_center(object_index), 211 | "top": bounding_box_data.get_top(object_index), 212 | "bottom": bounding_box_data.get_bottom(object_index), 213 | "back": bounding_box_data.get_back(object_index), 214 | "front": bounding_box_data.get_front(object_index), 215 | "left": bounding_box_data.get_left(object_index), 216 | "right": bounding_box_data.get_right(object_index) 217 | } 218 | 219 | self.agent_state_data.append(state_data) 220 | self.idx += 1 221 | 222 | def stop_agent(self): 223 | self.agent_state["velocity"] = [0, 0, 0] 224 | return self.tdw_object.communicate( 225 | [{"$type": "set_object_drag", "id": self.agent_id, "drag": 100, "angular_drag": 100}, 226 | {"$type": "step_physics", "frames": 1}, 227 | {"$type": "set_object_drag", "id": self.agent_id, "drag": 0, "angular_drag": 100} 228 | ]) 229 | 230 | def approach(self, till_x=None, till_y=None, till_z=None, approach_dist=0.01, speed_threshold=1.3, rotation=0): 231 | speed = 0 232 | force_magnitude = 0.8 233 | distance_to_travel = 10000 234 | rotated_angle = 0 235 | unit_rotation = 5 236 | 237 | if till_x is not None and till_y is not None and till_z is not None: 238 | total_force = [0, 0, 0] 239 | while True: 240 | # Find the required velocity 241 | # vect_to_target = np.array((-self.agent_position["x"] + till_x, -self.agent_position["y"] + till_y, 242 | # -self.agent_position["z"] + till_z)) 243 | # v_hat_ = vect_to_target / np.linalg.norm(vect_to_target) 244 | # 245 | # v_hat = [v_hat_[0] * speed_threshold - self.agent_state["velocity"][0], 246 | # v_hat_[1] * speed_threshold - self.agent_state["velocity"][1], 247 | # v_hat_[2] * speed_threshold - self.agent_state["velocity"][2]] 248 | # force_applied = {"x": v_hat[0], "y": v_hat[1], "z": v_hat[2]} 249 | # # Calculate Rotation 250 | # direction = 1 if rotation - rotated_angle > 0 else -1 251 | # angle = direction * min(abs(rotation - rotated_angle), unit_rotation) 252 | # rotated_angle += angle 253 | # self.agent_angle += angle 254 | # qo = euler_to_quaternion(0, convert_to_radians(self.agent_angle), 0) 255 | # 256 | # total_force = [e + abs(force_applied[f]) for e, f in zip(total_force, ["x", "y", "z"])] 257 | # self.forces.append(force_applied.copy()) 258 | # self.total_force = [e + abs(force_applied[f]) for e, f in zip(self.total_force, ["x", "y", "z"])] 259 | # resp = self.tdw_object.communicate( 260 | # [{"$type": "apply_force_to_object", "force": force_applied.copy(), 261 | # "id": self.agent_id}, 262 | # {"$type": "rotate_object_to", "rotation": {"w": qo[3], "x": qo[0], "y": qo[1], "z": qo[2]}, 263 | # "id": self.agent_id} 264 | # ] 265 | # ) 266 | # self.save_imag(resp) 267 | # velocity, _ = self.get_agent_loc_vel(resp, both=True) 268 | # 269 | # distance_to_travel = math.sqrt( 270 | # (till_x - self.agent_position["x"]) ** 2 + (till_y - self.agent_position["y"]) ** 2 271 | # + (till_z - self.agent_position["z"]) ** 2 272 | # ) 273 | # new_vect = np.array((-self.agent_position["x"] + till_x, -self.agent_position["y"] + till_y, 274 | # -self.agent_position["z"] + till_z)) 275 | # new_vect = new_vect / np.linalg.norm(new_vect) 276 | # if distance_to_travel < approach_dist or np.dot(new_vect, v_hat_) < 0: 277 | # break 278 | # Calculate Rotation 279 | # direction = 1 if rotation - rotated_angle > 0 else -1 280 | # angle = direction * min(abs(rotation - rotated_angle), unit_rotation) 281 | # rotated_angle += angle 282 | # self.agent_angle += angle 283 | # qo = euler_to_quaternion(0, convert_to_radians(self.agent_angle), 0) 284 | if speed < speed_threshold: 285 | vect_to_target = np.array((-self.agent_position["x"] + till_x, 286 | -self.agent_position["y"] + till_y, 287 | -self.agent_position["z"] + till_z)) 288 | v_hat = vect_to_target / np.linalg.norm(vect_to_target) 289 | v_hat_ = [v_hat[0] * speed_threshold, v_hat[1] * speed_threshold, v_hat[2] * speed_threshold] 290 | v_hat = [v_hat_[0] * speed_threshold - self.agent_state["velocity"][0], 291 | v_hat_[1] * speed_threshold - self.agent_state["velocity"][1], 292 | v_hat_[2] * speed_threshold - self.agent_state["velocity"][2], 293 | ] 294 | force_applied = {"x": v_hat[0], "y": v_hat[1], "z": v_hat[2]} 295 | self.total_force = [e + abs(force_applied[f]) for e, f in 296 | zip(self.total_force, ["x", "y", "z"])] 297 | resp = self.tdw_object.communicate( 298 | [{"$type": "apply_force_to_object", 299 | "force": force_applied.copy(), 300 | "id": self.agent_id} 301 | # {"$type": "rotate_object_to", "rotation": {"w": qo[3], "x": qo[0], "y": qo[1], "z": qo[2]}, 302 | # "id": self.agent_id} 303 | ] 304 | ) 305 | else: 306 | resp = self.tdw_object.communicate([{"$type": "do_nothing"} 307 | ]) 308 | force_applied = {"x": 0, "y": 0.0, "z": 0} 309 | 310 | velocity, _ = self.get_agent_loc_vel(resp, both=True) 311 | velocity_x, velocity_y, velocity_z = velocity[0], velocity[1], velocity[2] 312 | speed = math.sqrt(velocity_x ** 2 + velocity_y ** 2 + velocity_z ** 2) 313 | previous_distance = distance_to_travel 314 | distance_to_travel = math.sqrt( 315 | (till_x - self.agent_position["x"]) ** 2 + (till_z - self.agent_position["z"]) ** 2 316 | ) 317 | if distance_to_travel < approach_dist or previous_distance < distance_to_travel: 318 | resp = self.stop_agent() 319 | self.forces.append({"x": -1 * velocity_x, "y": 0.0, "z": -1 * velocity_z}) 320 | self.save_imag(resp) 321 | break 322 | else: 323 | self.save_imag(resp) 324 | self.forces.append(force_applied.copy()) 325 | 326 | elif till_x is not None and till_z is not None: 327 | 328 | total_force = [0, 0, 0] 329 | while True: 330 | # Find the required velocity 331 | vect_to_target = np.array((-self.agent_position["x"] + till_x, -self.agent_position["z"] + till_z)) 332 | v_hat_ = vect_to_target / np.linalg.norm(vect_to_target) 333 | v_hat = [v_hat_[0]*speed_threshold - self.agent_state["velocity"][0], v_hat_[1]*speed_threshold - self.agent_state["velocity"][2]] 334 | force_applied = {"x": v_hat[0], "y": 0.0, "z": v_hat[1]} 335 | # Calculate Rotation 336 | direction = 1 if rotation - rotated_angle > 0 else -1 337 | angle = direction * min(abs(rotation - rotated_angle), unit_rotation) 338 | rotated_angle += angle 339 | self.agent_angle += angle 340 | qo = euler_to_quaternion(0, convert_to_radians(self.agent_angle), 0) 341 | 342 | total_force = [e + abs(force_applied[f]) for e, f in zip(total_force, ["x", "y", "z"])] 343 | self.forces.append(force_applied.copy()) 344 | self.total_force = [e + abs(force_applied[f]) for e, f in zip(self.total_force, ["x", "y", "z"])] 345 | resp = self.tdw_object.communicate( 346 | [{"$type": "apply_force_to_object", "force": force_applied.copy(), 347 | "id": self.agent_id}, 348 | {"$type": "rotate_object_to", "rotation": {"w": qo[3], "x": qo[0], "y": qo[1], "z": qo[2]}, 349 | "id": self.agent_id} 350 | ] 351 | ) 352 | self.save_imag(resp) 353 | velocity, _ = self.get_agent_loc_vel(resp, both=True) 354 | 355 | distance_to_travel = math.sqrt( 356 | (till_x - self.agent_position["x"])**2 + (till_z - self.agent_position["z"])**2 357 | ) 358 | new_vect = np.array((-self.agent_position["x"] + till_x, -self.agent_position["z"] + till_z)) 359 | new_vect = new_vect/np.linalg.norm(new_vect) 360 | if distance_to_travel < approach_dist or np.dot(new_vect, v_hat_) < 0: 361 | break 362 | elif till_x is not None: 363 | force = self.direction * -force_magnitude 364 | while True: 365 | if abs(speed) < speed_threshold: 366 | self.total_force[0] += force_magnitude 367 | self.forces.append({"x": force, "y": 0.0, "z": 0}) 368 | resp = self.tdw_object.communicate( 369 | {"$type": "apply_force_to_object", "force": {"x": force, "y": 0.0, "z": 0}, "id": self.agent_id}) 370 | else: 371 | resp = self.tdw_object.communicate({"$type": "do_nothing"}) 372 | self.forces.append({"x": 0, "y": 0.0, "z": 0}) 373 | speed, _ = self.get_agent_loc_vel(resp) 374 | self.save_imag(resp) 375 | distance_to_travel = abs(till_x - self.agent_position["x"]) 376 | if distance_to_travel < 0.01: 377 | resp = self.tdw_object.communicate([{"$type": "set_object_drag", "id": self.agent_id, "drag": 100, "angular_drag": 100}, 378 | {"$type": "step_physics", "frames": 1}, 379 | {"$type": "set_object_drag", "id": self.agent_id, "drag": 0, "angular_drag": 100}, 380 | {"$type": "teleport_object", "position": {"x": till_x, "y": self.agent_position["y"], "z": self.agent_position["z"]}, "id": self.agent_id}]) 381 | self.save_imag(resp) 382 | self.forces.append({"x": -1*speed, "y": 0.0, "z": 0}) 383 | break 384 | 385 | elif till_z is not None: 386 | force = force_magnitude if till_z > self.agent_position["z"] else -force_magnitude 387 | while True: 388 | if abs(speed) < speed_threshold: 389 | self.total_force[2] += force_magnitude 390 | self.forces.append({"x": 0.0, "y": 0.0, "z": force}) 391 | resp = self.tdw_object.communicate( 392 | {"$type": "apply_force_to_object", "force": {"x": 0.0, "y": 0.0, "z": force}, "id": self.agent_id}) 393 | else: 394 | resp = self.tdw_object.communicate({"$type": "do_nothing"}) 395 | self.forces.append({"x": 0.0, "y": 0.0, "z": 0.0}) 396 | speed, _ = self.get_agent_loc_vel(resp, get_x=False) 397 | self.save_imag(resp) 398 | distance_to_travel = abs(till_z - self.agent_position["z"]) 399 | if distance_to_travel < 0.01: 400 | resp = self.tdw_object.communicate([{"$type": "set_object_drag", "id": self.agent_id, "drag": 100, "angular_drag": 100}, 401 | {"$type": "step_physics", "frames": 1}, 402 | {"$type": "set_object_drag", "id": self.agent_id, "drag": 0, "angular_drag": 100}, 403 | {"$type": "teleport_object", 404 | "position": {"x": self.agent_position["x"], 405 | "y": self.agent_position["y"], 406 | "z": till_z}, 407 | "id": self.agent_id}]) 408 | self.forces.append({"x": 0, "y": 0.0, "z": -1*speed}) 409 | self.save_imag(resp) 410 | break 411 | 412 | def settle_drop_object(self, n=100, stop_at_base=False, ground_height=0.11, residual_force=None): 413 | for i in range(n): 414 | resp = self.tdw_object.communicate({"$type": "do_nothing"}) 415 | self.save_imag(resp) 416 | if i == 0 and residual_force is not None: 417 | self.forces.append(residual_force) 418 | else: 419 | self.forces.append({"x": 0, "y": 0, "z": 0}) 420 | speed = np.linalg.norm(self.agent_state["velocity"]) 421 | if speed < 0.01: 422 | return 423 | if self.agent_position["y"] <= 1.01*ground_height and stop_at_base and i > 10: 424 | # breaking_force = [abs(e) for e in self.agent_state["velocity"]] 425 | # resp = self.stop_agent() 426 | # self.save_imag(resp) 427 | # self.forces.append({"x": breaking_force[0], "y": 0, "z": breaking_force[2]}) 428 | # print(i) 429 | # exit() 430 | # # self.forces[-1] = {"x": breaking_force[0], "y": 0, "z": breaking_force[2]} 431 | return 432 | 433 | def jump(self, height, width, direction_x_z=(1, 0), hight_ground_height=None): 434 | 435 | breaking_force = [abs(e) for e in self.agent_state["velocity"]] 436 | self.forces.append({"x": breaking_force[0], "y": breaking_force[1], "z": breaking_force[2]}) 437 | self.save_imag(self.stop_agent()) 438 | 439 | ground_height = self.agent_position["y"] if hight_ground_height is None else hight_ground_height 440 | jump_forces_list_obj = jump_forces_list(direction=self.direction, direction_x_z=direction_x_z) 441 | force_vect = jump_forces_list_obj.get_jump_force(self.agent_state["velocity"], height, width) 442 | 443 | resp = self.tdw_object.communicate( 444 | {"$type": "apply_force_to_object", "force": {"x": force_vect[0], "y": force_vect[1], "z": force_vect[2]}, "id": self.agent_id}) 445 | self.forces.append({"x": force_vect[0], "y": force_vect[1], "z": force_vect[2]}) 446 | self.save_imag(resp) 447 | 448 | self.settle_drop_object(ground_height=ground_height, stop_at_base=True, n=200) 449 | self.total_force = [e + abs(f) for e, f in zip(self.total_force, force_vect)] 450 | 451 | 452 | def reveal_occluder(self, occluder): 453 | resp = self.tdw_object.communicate( 454 | {"$type": "apply_force", "origin": {"x": occluder.occluder_position["x"] - 0.5 * occluder.reveal_normal[0], 455 | "y": 0.55, 456 | "z": occluder.occluder_position["z"] - 0.5 * occluder.reveal_normal[2]}, 457 | "target": {"x": occluder.occluder_position["x"] + 0.5 * occluder.reveal_normal[0], 458 | "y": 0.55, "z": occluder.occluder_position["z"] + 0.5 * occluder.reveal_normal[2]}, 459 | "magnitude": 2} 460 | ) 461 | self.save_imag(resp) 462 | for i in range(92): 463 | resp = self.tdw_object.communicate({"$type": "do_nothing"}) 464 | -------------------------------------------------------------------------------- /util/pit.py: -------------------------------------------------------------------------------- 1 | from util import utils 2 | 3 | 4 | class Pit: 5 | def __init__(self, config_file): 6 | self.config_file = config_file 7 | self.pit_width = 0 8 | self.pit_ids = [] 9 | self.pit_widths = [] 10 | self.bridge = False 11 | self.top_pt, self.bottom_pt = None, None 12 | self.bridge_1, self.bridge_2 = None, None 13 | self.scene_3_4 = False 14 | def create_pit(self, tdw_object, scene_config, scene_3_4=False): 15 | self.scene_3_4 = scene_3_4 16 | if not scene_3_4: 17 | if self.config_file["pit_width"] != -1: 18 | self.pit_width = scene_config.bridge_width[self.config_file["pit_width"]] 19 | pit_side_1 = scene_config.pit_positions[0][0].copy() 20 | pit_side_2 = scene_config.pit_positions[0][1].copy() 21 | self.pit_side_width = abs(pit_side_1["x"] - scene_config.positions_x_scene_1_2[1]) 22 | self.pit_side_width -= (self.pit_width/2) 23 | self.pit_side_width *= 2 24 | if self.config_file["obstacle_pos_x"] == -1: 25 | pit_side_1["x"] -= (1.2 + 0.16 + self.pit_width) 26 | pit_side_2["x"] -= (1.2 + 0.16 + self.pit_width) 27 | self.pit_ids.append(utils.create_raised_ground(tdw_object, pit_side_1, 28 | width=self.pit_side_width)) 29 | self.pit_widths.append(self.pit_side_width) 30 | self.pit_ids.append(utils.create_raised_ground(tdw_object, pit_side_2, 31 | width=self.pit_side_width)) 32 | self.pit_widths.append(self.pit_side_width) 33 | else: 34 | self.pit_ids.append(utils.create_raised_ground(tdw_object, scene_config.pit_positions[0][1], 35 | width=12)) 36 | self.pit_widths.append(12) 37 | if self.config_file["pit_depth"] != -1: 38 | side_1_pos = scene_config.pit_positions[2][0].copy() 39 | side_2_pos = scene_config.pit_positions[2][1].copy() 40 | if self.config_file["obstacle_pos_x"] == -1: 41 | side_1_pos["x"] -= (1.2 + 0.16 + self.pit_width) 42 | side_2_pos["x"] -= (1.2 + 0.16 + self.pit_width) 43 | side_1_pos["z"] = scene_config.positions_z_scene_1_2[self.config_file["obstacle_pos_z"]] + 2.5 44 | side_2_pos["z"] = scene_config.positions_z_scene_1_2[self.config_file["obstacle_pos_z"]] - 2.5 45 | dist_btw_pts = abs(side_1_pos["z"] - side_2_pos["z"]) 46 | sides_depth = dist_btw_pts - scene_config.pit_depth[self.config_file["pit_depth"]] 47 | side_width = scene_config.bridge_width[self.config_file["pit_width"]] 48 | self.bottom_pt = side_1_pos.copy() 49 | self.bottom_pt["z"] -= (sides_depth/2 - 0.2) 50 | self.top_pt = side_2_pos.copy() 51 | self.top_pt["z"] += (sides_depth / 2 - 0.2) 52 | self.pit_ids.append(utils.create_raised_ground(tdw_object, side_1_pos, 53 | width=side_width, depth=sides_depth)) 54 | self.pit_widths.append(side_width) 55 | self.pit_ids.append(utils.create_raised_ground(tdw_object, side_2_pos, 56 | width=side_width, depth=sides_depth)) 57 | self.pit_widths.append(side_width) 58 | else: 59 | assert not (self.config_file["pit_width_1"] == -1 and self.config_file[ 60 | "pit_width_2"] == -1), "Don't use pit " \ 61 | "barrier when you" \ 62 | " don't intend " \ 63 | "to have a pit" 64 | pit_side_1 = scene_config.pit_positions[1][0].copy() 65 | pit_side_2 = scene_config.pit_positions[1][2].copy() 66 | if self.config_file["pit_width_1"] != -1: 67 | dist_1 = abs(pit_side_1["x"] - scene_config.positions[1]["x"]) 68 | self.pit_width_1 = scene_config.bridge_width_3_4[self.config_file["pit_width_1"]] 69 | self.side_1_width = (dist_1 - self.pit_width_1/2)*2 70 | else: 71 | dist_1 = abs(pit_side_1["x"] - scene_config.positions[3]["x"]) 72 | self.side_1_width = (dist_1 - scene_config.bridge_width_3_4[self.config_file["pit_width_2"]] / 2) * 2 73 | if self.config_file["pit_width_2"] != -1: 74 | dist_2 = abs(pit_side_2["x"] - scene_config.positions[3]["x"]) 75 | self.pit_width_2 = scene_config.bridge_width_3_4[self.config_file["pit_width_2"]] 76 | self.side_2_width = (dist_2 - self.pit_width_2 / 2)*2 77 | else: 78 | dist_2 = abs(pit_side_2["x"] - scene_config.positions[1]["x"]) 79 | self.side_2_width = (dist_2 - scene_config.bridge_width_3_4[self.config_file["pit_width_1"]] / 2) * 2 80 | 81 | self.pit_ids.append(utils.create_raised_ground(tdw_object, pit_side_1, 82 | width=self.side_1_width)) 83 | self.pit_widths.append(self.side_1_width) 84 | self.pit_ids.append(utils.create_raised_ground(tdw_object, pit_side_2, 85 | width=self.side_2_width)) 86 | self.pit_widths.append(self.side_2_width) 87 | 88 | if self.config_file["pit_width_1"] != -1 and self.config_file["pit_width_2"] != -1: 89 | self.side_mid_width = abs(scene_config.positions[1]["x"] - scene_config.positions[3]["x"]) - \ 90 | (self.pit_width_1/2 + self.pit_width_2/2) 91 | 92 | pit_mid = pit_side_1.copy() 93 | pit_mid["x"] = scene_config.positions[1]["x"] - (self.side_mid_width/2 + self.pit_width_1/2) 94 | self.pit_ids.append(utils.create_raised_ground(tdw_object, pit_mid, 95 | width=self.side_mid_width)) 96 | self.pit_widths.append(self.side_mid_width) 97 | if self.config_file["barrier_type"] == "pit-with-bridge": 98 | self.bridge = True 99 | if "bridge_1_z" in self.config_file: 100 | if self.config_file["bridge_1_z"] != -1: 101 | self.bridge_1_position = pit_side_1.copy() 102 | self.bridge_1_position["x"] = scene_config.positions[1]["x"] 103 | self.bridge_1_position["z"] = scene_config.positions_z[self.config_file["bridge_1_z"]] 104 | self.pit_ids.append(utils.create_raised_ground(tdw_object, self.bridge_1_position, 105 | width=self.pit_width_1, depth=0.6)) 106 | self.pit_widths.append(self.pit_width_1) 107 | self.bridge_1 = self.pit_ids[-1] 108 | if "bridge_2_z" in self.config_file: 109 | if self.config_file["bridge_2_z"] != -1: 110 | self.bridge_2_position = pit_side_1.copy() 111 | self.bridge_2_position["x"] = scene_config.positions[3]["x"] 112 | self.bridge_2_position["z"] = scene_config.positions_z[self.config_file["bridge_2_z"]] 113 | self.pit_ids.append(utils.create_raised_ground(tdw_object, self.bridge_2_position, 114 | width=self.pit_width_2, depth=0.6)) 115 | self.pit_widths.append(self.pit_width_2) 116 | self.bridge_2 = self.pit_ids[-1] 117 | 118 | # self.pit_ids.append(utils.create_raised_ground(tdw_object, {"x": 1.009, "y":0.462, "z":-6.34}, 119 | # width=1.26, depth=2)) 120 | # self.pit_ids.append(utils.create_raised_ground(tdw_object, {"x": -0.69, "y": 0.462, "z": -6.34}, 121 | # width=1.26, depth=2)) 122 | # self.pit_ids.append(utils.create_raised_ground(tdw_object, {"x": 1.009, "y": 0.462, "z": 1.31}, 123 | # width=1.26, depth=7)) 124 | # self.pit_ids.append(utils.create_raised_ground(tdw_object, {"x": -0.69, "y": 0.462, "z": 1.31}, 125 | # width=1.26, depth=7)) 126 | 127 | 128 | 129 | if self.config_file["barrier_type"] == "pit-with-bridge": 130 | if not scene_3_4: 131 | self.bridge = True 132 | self.bridge_position = scene_config.bridge_position[0] 133 | self.bridge_width = scene_config.bridge_width[self.config_file["pit_width"]] 134 | self.bridge_position["z"] = scene_config.positions_z_scene_1_2[self.config_file["bridge_z"]] 135 | self.pit_ids.append(utils.create_raised_ground(tdw_object, self.bridge_position, 136 | width=self.bridge_width, 137 | depth=0.6)) 138 | self.pit_widths.append(self.bridge_width) 139 | self.bridge_1 = self.pit_ids[-1] 140 | return self.pit_ids 141 | -------------------------------------------------------------------------------- /util/ramp.py: -------------------------------------------------------------------------------- 1 | 2 | class Ramp: 3 | def __init__(self, position, rotation, height, agent_start_position): 4 | self.position = position 5 | self.rotation = rotation 6 | self.ramp_start = position 7 | self.height = height 8 | self.jump_point_1 = position 9 | self.jump_point_2 = position 10 | self.jump_point_3 = position 11 | self.agent_start_position = agent_start_position 12 | self.ramp_base_id, self.ramp_slope_id = [None]*2 13 | 14 | def get_positions(self, platform_only): 15 | self.platform_side = (0, 1) if self.rotation in [90, -90] else (1, 0) 16 | self.direction = 1 if self.rotation in [90, 180] else -1 17 | if self.height == 0.8: 18 | slope_position = {"x": self.position["x"] + self.platform_side[0]*self.direction*0.268, 19 | "y": 0.388, 20 | "z": self.position["z"] + self.platform_side[1]*self.direction*0.268} 21 | slope_rotation = {"x": 0, "y": self.rotation, "z": -36.342} 22 | slope_scale = {"x": 0.05, "y": 1.0, "z": 0.7} 23 | if platform_only: 24 | base_position = self.position.copy() 25 | base_position["y"] = 0.409 26 | else: 27 | base_position = {"x": self.position["x"] - self.platform_side[0]*self.direction*0.358, 28 | "y": 0.409, "z": self.position["z"] - self.platform_side[1]*self.direction*-0.358} 29 | base_scale = {"x": 0.7, "y": 0.8, "z": 0.7} 30 | self.ramp_start = {"x": slope_position["x"] + self.platform_side[0] * self.direction * 0.565, "y": 0, 31 | "z": slope_position["z"] + self.platform_side[1] * self.direction * 0.565} 32 | self.ramp_start_ = {"x": slope_position["x"] + self.platform_side[0] * self.direction * 0.4, "y": 0, 33 | "z": slope_position["z"] + self.platform_side[1] * self.direction * 0.4} 34 | 35 | # self.jump_point_1 = {"x": base_position["x"] - self.platform_side[0] * self.direction * 0.387, "y": 0, 36 | # "z": base_position["z"] - self.platform_side[1] * self.direction * 0.387} 37 | self.target_position = base_position.copy() 38 | self.target_position["y"] = 0.905 39 | elif self.height == 0.5: 40 | slope_position = {"x": self.position["x"] + self.platform_side[0]*self.direction*0.418, 41 | "y": 0.229, 42 | "z": self.position["z"] + self.platform_side[1]*self.direction*0.418} 43 | slope_rotation = {"x": 0, "y": self.rotation, "z": -60} 44 | slope_scale = {"x": 0.05, "y": 1.0, "z": 0.7} 45 | if platform_only: 46 | base_position = self.position.copy() 47 | base_position["y"] = 0.25 48 | else: 49 | base_position = {"x": self.position["x"] - self.platform_side[0]*self.direction*0.358, 50 | "y": 0.25, "z": self.position["z"] - self.platform_side[1]*self.direction*0.358} 51 | base_scale = {"x": 0.7, "y": 0.5, "z": 0.7} 52 | self.ramp_start = {"x": slope_position["x"] + self.platform_side[0] * self.direction * 0.565, "y": 0, 53 | "z": slope_position["z"] + self.platform_side[1] * self.direction * 0.565} 54 | self.ramp_start_ = {"x": slope_position["x"] + self.platform_side[0] * self.direction * 0.42, "y": 0, 55 | "z": slope_position["z"] + self.platform_side[1] * self.direction * 0.42} 56 | # self.jump_point_1 = {"x": base_position["x"] - self.platform_side[0] * self.direction * 0.387, "y": 0, 57 | # "z": base_position["z"] - self.platform_side[1] * self.direction * 0.387} 58 | self.target_position = base_position.copy() 59 | self.target_position["y"] = 0.6 60 | elif self.height == 0.2: 61 | slope_position = {"x": self.position["x"] + self.platform_side[0]*self.direction*0.474, 62 | "y": 0.077, 63 | "z": self.position["z"] + self.platform_side[1]*self.direction*0.474} 64 | slope_rotation = {"x": 0.0, "y": self.rotation, "z": -77.979} 65 | slope_scale = {"x": 0.05, "y": 1.0, "z": 0.7} 66 | if platform_only: 67 | base_position = self.position.copy() 68 | base_position["y"] = 0.103 69 | else: 70 | base_position = {"x": self.position["x"] - self.platform_side[0]*self.direction*0.358, 71 | "y": 0.103, "z": self.position["z"] - self.platform_side[1]*self.direction*0.358} 72 | base_scale = {"x": 0.7, "y": 0.2, "z": 0.7} 73 | self.ramp_start = {"x": slope_position["x"] + self.platform_side[0]*self.direction*0.738, "y": 0, 74 | "z": slope_position["z"] + self.platform_side[1]*self.direction*0.738} 75 | self.ramp_start_ = {"x": slope_position["x"] + self.platform_side[0] * self.direction * 0.6, "y": 0, 76 | "z": slope_position["z"] + self.platform_side[1] * self.direction * 0.6} 77 | # self.jump_point_1 = {"x": base_position["x"] - self.platform_side[0] * self.direction * 0.387, "y": 0, 78 | # "z": base_position["z"] - self.platform_side[1] * self.direction * 0.387} 79 | self.target_position = base_position.copy() 80 | self.target_position["y"] = 0.303 81 | 82 | self.calculate_platform_points(base_position) 83 | self.positions = { 84 | "slope_position": slope_position, 85 | "slope_rotation": slope_rotation, 86 | "slope_scale": slope_scale, 87 | "base_position" : base_position, 88 | "base_rotation": {"x": 0.0, "y": 0.0, "z": 0.0}, 89 | "base_scale": base_scale 90 | } 91 | return self.positions 92 | 93 | def calculate_platform_points(self, base_position): 94 | offset = 0.76 95 | offset_ = 0.898 96 | self.jump_point_1 = {"x": base_position["x"] - self.platform_side[0] * self.direction * offset, "y": 0, 97 | "z": base_position["z"] - self.platform_side[1] * self.direction * offset} 98 | self.jump_point_1_ = {"x": base_position["x"] - self.platform_side[0] * self.direction * offset_, "y": 0, 99 | "z": base_position["z"] - self.platform_side[1] * self.direction * offset} 100 | self.jump_point_2 = {"x": base_position["x"] + self.platform_side[1] * self.direction * offset, "y": 0, 101 | "z": base_position["z"] + self.platform_side[0] * self.direction * offset} 102 | self.jump_point_2_ = {"x": base_position["x"] + self.platform_side[1] * self.direction * offset_, "y": 0, 103 | "z": base_position["z"] + self.platform_side[0] * self.direction * offset} 104 | self.jump_point_3 = {"x": base_position["x"] - self.platform_side[1] * self.direction * offset, "y": 0, 105 | "z": base_position["z"] - self.platform_side[0] * self.direction * offset} 106 | self.jump_point_3_ = {"x": base_position["x"] - self.platform_side[1] * self.direction * offset_, "y": 0, 107 | "z": base_position["z"] - self.platform_side[0] * self.direction * offset} 108 | -------------------------------------------------------------------------------- /util/scene_configuration.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | class SceneConfiguration: 4 | def __init__(self, scene_offset=(0, 0, -3.806)): 5 | # self.scene_offset = (0, 0, -4.026) 6 | # self.scene_offset = (0, 0, -4.032) 7 | self.pit_depth_ = 0.8 8 | self.scene_offset = scene_offset 9 | self.agent_colors = [ 10 | {"r": 240 / 255, "g": 19 / 255, "b": 77 / 255, "a": 1.0}, 11 | {"r": 7 / 255, "g": 121 / 255, "b": 228 / 255, "a": 1.0}, 12 | {"r": 52 / 255, "g": 46 / 255, "b": 173 / 255, "a": 1.0}, 13 | {"r": 253 / 255, "g": 46 / 255, "b": 179 / 255, "a": 1.0}, 14 | {"r": 237 / 255, "g": 12 / 255, "b": 239 / 255, "a": 1.0}, 15 | ] 16 | # self.target_colors = [ 17 | # 18 | # {"r": 149 / 255, "g": 56 / 255, "b": 158 / 255, "a": 1.0}, 19 | # {"r": 216 / 255, "g": 52 / 255, "b": 95 / 255, "a": 1.0}, 20 | # {"r": 148 / 255, "g": 252 / 255, "b": 19 / 255, "a": 1.0}, 21 | # {"r": 255 / 255, "g": 87 / 255, "b": 34 / 255, "a": 1.0}, 22 | # {"r": 210 / 255, "g": 230 / 255, "b": 3 / 255, "a": 1.0}, 23 | # {"r": 244 / 255, "g": 89 / 255, "b": 5 / 255, "a": 1.0}, 24 | # {"r": 243 / 255, "g": 14 / 255, "b": 92 / 255, "a": 1.0}, 25 | # 26 | # {"r": 67 / 255, "g": 216 / 255, "b": 210 / 255, "a": 1.0}, 27 | # {"r": 88 / 255, "g": 141 / 255, "b": 168 / 255, "a": 1.0}, 28 | # {"r": 48 / 255, "g": 71 / 255, "b": 94 / 255, "a": 1.0}, 29 | # {"r": 0 / 255, "g": 188 / 255, "b": 212 / 255, "a": 1.0}, 30 | # {"r": 217 / 255, "g": 191 / 255, "b": 119 / 255, "a": 1.0}, 31 | # {"r": 85 / 255, "g": 34 / 255, "b": 68 / 255, "a": 1.0}, 32 | # {"r": 8 / 255, "g": 255 / 255, "b": 200 / 255, "a": 1.0}, 33 | # 34 | # ] 35 | self.target_colors_ = { 36 | "set_1": [ 37 | {"r": 216 / 255, "g": 52 / 255, "b": 95 / 255, "a": 1.0}, 38 | {"r": 148 / 255, "g": 252 / 255, "b": 19 / 255, "a": 1.0}, 39 | {"r": 255 / 255, "g": 87 / 255, "b": 34 / 255, "a": 1.0}, 40 | {"r": 210 / 255, "g": 230 / 255, "b": 3 / 255, "a": 1.0}, 41 | ], 42 | "set_2": [ 43 | {"r": 88 / 255, "g": 141 / 255, "b": 168 / 255, "a": 1.0}, 44 | {"r": 48 / 255, "g": 71 / 255, "b": 94 / 255, "a": 1.0}, 45 | {"r": 0 / 255, "g": 188 / 255, "b": 212 / 255, "a": 1.0}, 46 | {"r": 85 / 255, "g": 34 / 255, "b": 68 / 255, "a": 1.0}, 47 | ] 48 | } 49 | 50 | self.target_colors = self.target_colors_["set_1"] + self.target_colors_["set_2"] 51 | 52 | self.ramp_position = [ 53 | {"x": 0 + self.scene_offset[0] + 1.66, "y": 0, "z": 0 + self.scene_offset[2]}, 54 | {"x": 0 + self.scene_offset[0] - 1.66, "y": 0, "z": 0 + self.scene_offset[2]} 55 | ] 56 | self.positions_x = [ 57 | 0 + self.scene_offset[0] + 1.66, 58 | 0 + self.scene_offset[0] + 0.83, 59 | 0 + self.scene_offset[0], 60 | 0 + self.scene_offset[0] - 0.83, 61 | 0 + self.scene_offset[0] - 1.66 62 | ] 63 | self.positions_x_without_barrier = [ 64 | 0 + self.scene_offset[0] + 1.66, 65 | 0 + self.scene_offset[0] + 1.106, 66 | 0 + self.scene_offset[0] + 0.5533, 67 | 0 + self.scene_offset[0], 68 | 0 + self.scene_offset[0] - 0.5533, 69 | 0 + self.scene_offset[0] - 1.106, 70 | 0 + self.scene_offset[0] - 1.66 71 | ] 72 | # self.positions_x = [ 73 | # 0 + self.scene_offset[0] + 1.86, 74 | # 0 + self.scene_offset[0] + 0.93, 75 | # 0 + self.scene_offset[0], 76 | # 0 + self.scene_offset[0] - 0.93, 77 | # 0 + self.scene_offset[0] - 1.86 78 | # ] 79 | # self.target_objects = [ 80 | # "bowl", 81 | # "cylinder", 82 | # "dumbbell", 83 | # "octahedron", 84 | # "pentagon", 85 | # "pipe", 86 | # "platonic", 87 | # "pyramid", 88 | # "torus", 89 | # "triangular_prism" 90 | # ] 91 | self.target_objects = { 92 | "target_objects_1": [ 93 | "pyramid", 94 | "bowl" 95 | ], 96 | "target_objects_2":[ 97 | "cylinder", 98 | "octahedron" 99 | ] 100 | } 101 | self.target_objects_ = self.target_objects["target_objects_1"] + self.target_objects["target_objects_2"] 102 | self.positions_x_scene_1_2 = [ 103 | 0 + (self.scene_offset[0] - 0.12) + 1.2, 104 | 0 + (self.scene_offset[0] - 0.12), 105 | 0 + (self.scene_offset[0] - 0.12) - 1.2, 106 | ] 107 | self.positions_x_scene_1_2_extended = [ 108 | 0 + (self.scene_offset[0] - 0.12) + 1.8, 109 | 0 + (self.scene_offset[0] - 0.12) - 1.8, 110 | ] 111 | self.positions_z_scene_1_2 = [ 112 | (self.scene_offset[2] + 1.56) - 0.549, 113 | 0 + (self.scene_offset[2] + 1.34), 114 | (self.scene_offset[2] + 1.34) + 0.549 115 | ] 116 | self.ramp_rotation = [0, 90, 180, -90] 117 | self.positions_z = [ 118 | self.scene_offset[2] - 0.549, 119 | 0 + self.scene_offset[2], 120 | self.scene_offset[2] + 0.549 121 | ] 122 | self.positions_y = [ 123 | scene_offset[1] + 0.102, 124 | 0.052, 125 | 0.102, 126 | 0.152, 127 | 0.202, 128 | 0.302, 129 | 0.352, 130 | 0.458, 131 | 0.509 132 | ] 133 | self.positions_y_platform_cube = [ 134 | 0.102, 135 | 0.2, 136 | ] 137 | self.platform_cube_height = [ 138 | 0.2, 139 | 0.4 140 | ] 141 | self.obstacle_z = [ 142 | self.scene_offset[2] - 0.222, 143 | 0 + self.scene_offset[2], 144 | self.scene_offset[2] + 0.222 145 | ] 146 | # x_pos -> Width -> Height -> Depth 147 | 148 | 149 | # x_pos -> occ position 150 | self.occluder_position = { 151 | "occ_z_0": {'x': 0.683, 'y': 0.617018, 'z': -0.295}, 152 | "occ_z_1": {'x': 0.645, 'y': 0.617018, 'z': -0.302}, 153 | "occ_z_2": {'x': 0.645, 'y': 0.617018, 'z': -0.302} 154 | } 155 | 156 | self.positions = [ 157 | {"x": 0 + self.scene_offset[0] + 1.66, "y": scene_offset[1] + 0.0011, 158 | "z": 0 + self.scene_offset[2]}, 159 | {"x": 0 + self.scene_offset[0] + 0.83, "y": scene_offset[1] + 0.0011, 160 | "z": 0 + self.scene_offset[2]}, 161 | {"x": 0 + self.scene_offset[0], "y": scene_offset[1] + 0.0011, 162 | "z": 0 + self.scene_offset[2]}, 163 | {"x": 0 + self.scene_offset[0] - 0.83, "y": scene_offset[1] + 0.0011, 164 | "z": 0 + self.scene_offset[2]}, 165 | {"x": 0 + self.scene_offset[0] - 1.66, "y": scene_offset[1] + 0.0011, 166 | "z": 0 + self.scene_offset[2]} 167 | ] 168 | self.pit_positions = [ 169 | ( 170 | {"x": (self.scene_offset[0] - 0.12) + 3.19, "y": 0.462 + self.pit_depth_/2, "z": 0 + scene_offset[2] + 0.896}, 171 | {"x": (self.scene_offset[0] - 0.12) - 3.19, "y": 0.462 + self.pit_depth_/2, "z": 0 + scene_offset[2] + 0.896}, 172 | {"x": (self.scene_offset[0] - 0.12), "y": 0.462 + self.pit_depth_/2, "z": 0 + scene_offset[2] + 0.896} 173 | ), 174 | ( 175 | {"x": self.scene_offset[0] + 3.75, "y": 0.462 + self.pit_depth_/2, "z": 0 + scene_offset[2] }, 176 | {"x": self.scene_offset[0], "y": 0.462 + self.pit_depth_/2, "z": 0 + scene_offset[2]}, 177 | {"x": self.scene_offset[0] - 3.75, "y": 0.462 + self.pit_depth_/2, "z": 0 + scene_offset[2]} 178 | ), 179 | ( 180 | {"x": 0 + (self.scene_offset[0] - 0.12), "y": 0.462 + self.pit_depth_/2, "z": 0 + scene_offset[2] + 1.122}, 181 | {"x": 0 + (self.scene_offset[0] - 0.12), "y": 0.462 + self.pit_depth_/2, "z": 0 + scene_offset[2] - 1.122} 182 | ) 183 | ] 184 | # 3.806 185 | self.bridge_position = [ 186 | {"x": (self.scene_offset[0] - 0.12), "y": 0.462 + self.pit_depth_/2 , "z": 0 + scene_offset[2] + 0.896}, 187 | ] 188 | # self.pit_width = [3.95, 3.35, 2.75] 189 | self.pit_width = [6, 5.5, 5] 190 | self.pit_width_3_4 = [(3, 2), (3, 2), (3, 2)] 191 | self.pit_depth = [1.1, 1.6, 2.1] 192 | self.bridge_width = [0.38, 0.6, 0.88, 1.13, 1.38] 193 | self.bridge_width_3_4 = [0.3, 0.5, 0.7] 194 | self.obstacle_height = [0.1, 0.2, 0.3, 0.4, 0.6, 0.7, 0.9, 1.0] 195 | # self.obstacle_height_pool = [ 196 | # [0.1, 0.4, 0.7], 197 | # [0.2, 0.6], # [0.2, 0.4, 0.7], 198 | # [0.3, 0.7], 199 | # ] 200 | self.obstacle_height_pool = [ 201 | [0.1, 0.3, 0.6], 202 | [0.2, 0.6], 203 | [0.3, 0.6], 204 | ] 205 | self.obstacle_depth_pool = [ 206 | [0.7, 1.5, 2.1], 207 | [0.6, 1.5, 2.0], 208 | [0.6, 1.1, 1.7, 2.1] 209 | ] 210 | self.obstacle_width_pool = [ 211 | [0.1, 0.4], 212 | [0.2, 0.5], 213 | [0.1, 0.5], 214 | [0.2, 0.4], 215 | [0.1, 0.2] 216 | ] 217 | self.obstacle_width = [0.1, 0.2, 0.4, 0.5] 218 | self.obstacle_depth = [0.6, 0.7, 1.1, 1.6, 2.1, 2.4, 2.6, 2.8, 3.0, 4.0] 219 | self.agent_types = ["cube", "sphere", "cone"] 220 | self.ramp_height = [0.2, 0.5, 0.8] 221 | self.floor_material = ["ceramic_tiles_beige_tan", "ceramic_tiles_copper", "ceramic_tiles_floral_white", "ceramic_tiles_brown_tomato", "tiles_hexagon_white", "ceramic_tiles_golden_sand", "ceramic_tiles_grey", "ceramic_tiles_green"] 222 | self.wall_material = ["square_padded_wall", "cinderblock_wall", "church_wall_chamfered_cracks"] 223 | 224 | 225 | 226 | -------------------------------------------------------------------------------- /util/target_object.py: -------------------------------------------------------------------------------- 1 | class TargetObject: 2 | def __init__(self, position, sphere_id): 3 | self.target_position = position 4 | self.target_id = sphere_id -------------------------------------------------------------------------------- /util/twogoalscene.py: -------------------------------------------------------------------------------- 1 | from util import physics_agent 2 | import random 3 | from util import utils 4 | import shutil 5 | import os 6 | from util.scene_configuration import SceneConfiguration 7 | import numpy as np 8 | import json 9 | import time 10 | import traceback 11 | from util import trajectory_generation 12 | from util import create_scene 13 | 14 | 15 | class TwoGoalScene: 16 | def __init__(self, tdw_object, camera_obj, args): 17 | self.tdw_object = tdw_object 18 | self.args = args 19 | self.camera_obj = camera_obj 20 | self.scene_config = SceneConfiguration() 21 | self.scene_offset = self.scene_config.scene_offset 22 | self.agent_shape = "cone" 23 | self.agent_color = random.sample(self.scene_config.agent_colors, 1)[0] 24 | self.target_1_color, self.target_2_color = random.sample(self.scene_config.target_colors, 2) 25 | self.target_1_color = {"r": 48 / 255, "g": 71 / 255, "b": 94 / 255, "a": 1.0} 26 | self.target_2_color = {"r": 216 / 255, "g": 52 / 255, "b": 95 / 255, "a": 1.0} 27 | self.agent_color = {"r": 7 / 255, "g": 121 / 255, "b": 228 / 255, "a": 1.0} 28 | self.high_scene = False 29 | self.pit_obj, self.pit_ids = None, [] 30 | 31 | def go_to_goal(self): 32 | # Check if obstacle is between goal and agent 33 | h_offset = 0.15 34 | offset = 0.147 35 | paths = [] 36 | # Cross barrier 37 | if self.config_file["barrier_type"] in ["pit", "pit-with-bridge"]: 38 | return trajectory_generation.cross_pit_scene_3_4(self.agent, self.config_file, 39 | [self.target_1_position, self.target_2_position], self.pit_obj) 40 | # if self.config_file["agent_pos_x"] in [0, 1]: 41 | # jump_point_x = self.scene_config.pit_positions[0][0]["x"] - self.scene_config.pit_width[ 42 | # self.config_file["pit_width"]] / 2 + 0.1 43 | # else: 44 | # jump_point_x = self.scene_config.pit_positions[0][1]["x"] + self.scene_config.pit_width[ 45 | # self.config_file["pit_width"]] / 2 - 0.1 46 | # if self.config_file["obj_pos_z"] > self.config_file["agent_pos_z"] or \ 47 | # self.config_file["obj_pos_z"] < self.config_file["agent_pos_z"]: 48 | # jump_point_z = self.scene_config.positions_z[1] 49 | # else: 50 | # jump_point_z = self.scene_config.positions_z[self.config_file["agent_pos_z"]] 51 | # 52 | # self.agent.approach(till_x=jump_point_x, till_z=jump_point_z) 53 | # self.agent.jump(height=None, width=self.scene_config.pit_width[self.config_file["pit_width"]]) 54 | 55 | elif self.config_file["barrier_type"] == "cube": 56 | paths.extend(self.cross_barrier()) 57 | return paths 58 | elif self.config_file["barrier_type"] == "barrier_with_door": 59 | paths.extend(self.cross_barrier(with_door=True)) 60 | return paths 61 | # # Go to goal object 62 | # if "obj_ramp_height" in self.config_file: 63 | # if self.scene_config.ramp_height[self.config_file["obj_ramp_height"]] == 0.5: 64 | # self.agent.approach(till_x=self.selected_target["x"], 65 | # till_z=self.selected_target["z"] + 1.26) 66 | # self.agent.approach(till_z=self.selected_target["z"] + 0.18) 67 | # if self.scene_config.ramp_height[self.config_file["obj_ramp_height"]] == 0.2: 68 | # self.agent.approach(till_x=self.selected_target["x"], 69 | # till_z=self.selected_target["z"] + 1.32) 70 | # self.agent.approach(till_z=self.selected_target["z"] + 0.18) 71 | # else: 72 | # self.agent.approach(till_x=self.selected_target["x"], 73 | # till_z=self.selected_target["z"], approach_dist=0.2) 74 | elif self.config_file["barrier_type"] == "ramp": 75 | # Get path for each ramp 76 | for ramp_obj in [self.ramp_1, self.ramp_2]: 77 | restrict_paths = [1, 2, 3] 78 | # if self.agent.agent_start_position["x"] > ramp_obj.position["x"] and ramp_obj.rotation == 180: 79 | # restrict_paths = [1, 2, 3] 80 | # elif self.agent.agent_start_position["x"] > ramp_obj.position["x"] and ramp_obj.rotation == 0: 81 | # restrict_paths = [2, 3, "R"] 82 | # elif self.agent.agent_start_position["x"] < ramp_obj.position["x"] and ramp_obj.rotation == 0: 83 | # restrict_paths = [1, 2, 3] 84 | # elif self.agent.agent_start_position["x"] < ramp_obj.position["x"] and ramp_obj.rotation == 180: 85 | # restrict_paths = [2, 3, "R"] 86 | # elif self.agent.agent_start_position["x"] > ramp_obj.position["x"] and ramp_obj.rotation == 90: 87 | # restrict_paths = [1, 3] 88 | # elif self.agent.agent_start_position["x"] > ramp_obj.position["x"] and ramp_obj.rotation == -90: 89 | # restrict_paths = [1, 2] 90 | # elif self.agent.agent_start_position["x"] < ramp_obj.position["x"] and ramp_obj.rotation == 90: 91 | # restrict_paths = [1, 2] 92 | # elif self.agent.agent_start_position["x"] < ramp_obj.position["x"] and ramp_obj.rotation == -90: 93 | # restrict_paths = [1, 3] 94 | paths.extend(trajectory_generation.cross_ramp(ramp_obj, self.agent, self.config_file, restrict_paths=restrict_paths)) 95 | 96 | return paths 97 | elif self.config_file["barrier_type"] == "platform": 98 | for ramp_obj, target_pos in zip([self.ramp_1, self.ramp_2], [self.target_1_position, self.target_2_position]): 99 | paths.extend(trajectory_generation.cross_platform(ramp_obj, self.agent, target_pos)) 100 | return paths 101 | 102 | def cross_barrier(self, with_door=False): 103 | paths = [] 104 | h_offset = 0.15 105 | offset = 0.147 106 | # Jump over barrier 107 | for obstacle, target in zip([self.obstacle_1, self.obstacle_2], [self.target_1_position, self.target_2_position]): 108 | 109 | if with_door: 110 | path = trajectory_generation.cross_barrier_with_door(obstacle, target, self.agent.agent_start_position) 111 | else: 112 | path = trajectory_generation.cross_barrier(obstacle, target, self.agent.agent_start_position, 113 | self.scene_config, config_file=self.config_file, 114 | inefficient=False) 115 | paths.extend(path) 116 | 117 | return paths 118 | 119 | def execute_config(self, config_file, output_dir): 120 | self.agent = physics_agent.Agent(self.tdw_object, 1, config_file["dir_name"]) 121 | shapes = self.scene_config.target_objects["target_objects_1"] + self.scene_config.target_objects["target_objects_2"] 122 | create_scene_arguments = { 123 | "config_file": config_file, 124 | "tdw_object": self.tdw_object, 125 | "high_scene": self.high_scene, 126 | "camera_obj": self.camera_obj, 127 | "target_1_color": {"r": 88 / 255, "g": 141 / 255, "b": 168 / 255, "a": 1.0}, 128 | "target_2_color": {"r": 148 / 255, "g": 252 / 255, "b": 19 / 255, "a": 1.0}, 129 | "target_1_shape": "sphere", 130 | "target_2_shape": "sphere", 131 | "agent_color": self.agent_color, 132 | "agent_shape": self.agent_shape, 133 | "agent": self.agent, 134 | "enable_image": self.args.enable_image, 135 | "actually_create": True 136 | } 137 | self.config_file = config_file 138 | return_state = create_scene.create_based_config_2_goal(**create_scene_arguments) 139 | self.obstacle_1 = return_state["obstacle_1"] 140 | self.obstacle_2 = return_state["obstacle_2"] 141 | self.ramp_1 = return_state["ramp_1"] 142 | self.ramp_2 = return_state["ramp_2"] 143 | self.high_scene = return_state["high_scene"] 144 | self.target_1_position = return_state["target_1_position"] 145 | self.target_2_position = return_state["target_2_position"] 146 | self.target_1_id = return_state["target_1_id"] 147 | self.target_2_id = return_state["target_2_id"] 148 | self.pit_obj = return_state["pit_obj"] 149 | if self.pit_obj is not None: 150 | self.pit_ids = self.pit_obj.pit_ids 151 | try: 152 | utils.settle_objects(self.tdw_object, n=20) 153 | paths = self.go_to_goal() 154 | time_required = [] 155 | for p in paths: 156 | start_time = time.time() 157 | self.agent.output_dir = os.path.join(self.agent.output_dir, "images") 158 | utils.settle_objects(self.tdw_object, n=10) 159 | self.execute_path(p) 160 | time_required.append(time.time() - start_time) 161 | # print( 162 | # f"Took {time.time() - start_time} secs. Average time per trajectory {sum(time_required) / len(time_required)}") 163 | unique_str = config_file["dir_name"] + f"_path_no_{self.agent.path_no}" 164 | with open(os.path.join(output_dir, "time.txt"), "a") as fp: 165 | fp.write( 166 | f"{unique_str} took {time.time() - start_time} secs. Average time per trajectory {sum(time_required) / len(time_required)} \n") 167 | 168 | self.reset() 169 | except Exception as err: 170 | # with open("error.txt", "a") as fp: 171 | # fp.write(f"Error in configuration {output_dir}\n") 172 | traceback.print_tb(err.__traceback__) 173 | print(err) 174 | 175 | commands = [{"$type": "destroy_object", "id": self.agent.agent_id}, 176 | {"$type": "destroy_object", "id": self.target_1_id}, 177 | {"$type": "destroy_object", "id": self.target_2_id} 178 | ] 179 | if self.obstacle_1 is not None and self.obstacle_2 is not None: 180 | if config_file["barrier_type"] == "barrier_with_door": 181 | if self.obstacle_1.barrier_1 is not None: 182 | commands.append({"$type": "destroy_object", "id": self.obstacle_1.barrier_1}) 183 | commands.append({"$type": "destroy_object", "id": self.obstacle_1.barrier_2}) 184 | commands.append({"$type": "destroy_object", "id": self.obstacle_1.barrier_3}) 185 | if self.obstacle_2.barrier_1 is not None: 186 | commands.append({"$type": "destroy_object", "id": self.obstacle_2.barrier_1}) 187 | commands.append({"$type": "destroy_object", "id": self.obstacle_2.barrier_2}) 188 | commands.append({"$type": "destroy_object", "id": self.obstacle_2.barrier_3}) 189 | else: 190 | if self.obstacle_1.obstacle_id is not None: 191 | commands.append({"$type": "destroy_object", "id": self.obstacle_1.obstacle_id}) 192 | if self.obstacle_2.obstacle_id is not None: 193 | commands.append({"$type": "destroy_object", "id": self.obstacle_2.obstacle_id}) 194 | for pit_id in self.pit_ids: 195 | commands.append( 196 | {"$type": "destroy_object", "id": pit_id} 197 | ) 198 | if self.ramp_1 is not None: 199 | if config_file["barrier_type"] == "ramp": 200 | commands.extend([ 201 | {"$type": "destroy_object", "id": self.ramp_1.ramp_slope_id}, 202 | ]) 203 | commands.extend([ 204 | {"$type": "destroy_object", "id": self.ramp_1.ramp_base_id} 205 | ]) 206 | if self.ramp_2 is not None: 207 | if config_file["barrier_type"] == "ramp": 208 | commands.extend([ 209 | {"$type": "destroy_object", "id": self.ramp_2.ramp_slope_id} 210 | 211 | ]) 212 | commands.extend([ 213 | {"$type": "destroy_object", "id": self.ramp_2.ramp_base_id} 214 | ]) 215 | os.makedirs(self.config_file["dir_name"], exist_ok=True) 216 | with open(os.path.join(self.config_file["dir_name"], "scene_config.json"), "w") as fp: 217 | json.dump(self.config_file, fp) 218 | 219 | self.tdw_object.communicate(commands) 220 | self.ramp_1, self.ramp_2, self.obstacle_1, self.obstacle_2, self.pit_obj, self.pit_ids = None, None, None, None, None, [] 221 | self.agent = None 222 | 223 | def reset(self): 224 | self.agent.reset() 225 | utils.teleport(self.tdw_object, self.target_1_position, self.target_1_id) 226 | utils.teleport(self.tdw_object, self.target_2_position, self.target_2_id) 227 | 228 | def execute_path(self, path): 229 | meta_data = None 230 | p1 = [e for e in path if e[0] == "goto"] 231 | ramp_obj = None 232 | 233 | if p1: 234 | self.agent.rotate_by_angle(90) 235 | self.agent.settle_drop_object(10, stop_at_base=False) 236 | self.agent.rotate_by_angle(-90) 237 | self.agent.settle_drop_object(3, stop_at_base=False) 238 | self.agent.rotate_by_angle(-90) 239 | self.agent.settle_drop_object(10, stop_at_base=False) 240 | self.agent.rotate_by_angle(90) 241 | self.agent.settle_drop_object(3, stop_at_base=False) 242 | p1 = p1[0] 243 | 244 | if p1[1] > self.agent.agent_position["x"]: 245 | self.agent.direction = -1 246 | if self.config_file["barrier_type"] in ["ramp", "platform"]: 247 | ramp_obj = self.ramp_1 248 | # self.agent.rotate_by_angle(90) 249 | else: 250 | self.agent.direction = 1 251 | if self.config_file["barrier_type"] in ["ramp", "platform"]: 252 | ramp_obj = self.ramp_2 253 | # self.agent.rotate_by_angle(-90) 254 | 255 | for point in path: 256 | if point[0] == "goto": 257 | if len(point) == 6: 258 | velocity_threshold = point[5] 259 | self.agent.approach(till_x=point[1], 260 | till_y=point[2], 261 | till_z=point[3], approach_dist=point[4], speed_threshold=velocity_threshold) 262 | elif len(point) == 5: 263 | angle = self.agent.calculate_angle([point[1], point[2]]) 264 | velocity_threshold = point[4] 265 | self.agent.approach(till_x=point[1], 266 | till_z=point[2], approach_dist=point[3], speed_threshold=velocity_threshold, 267 | rotation=angle) 268 | else: 269 | angle = self.agent.calculate_angle([point[1], point[2]]) 270 | velocity_threshold = 1.3 271 | self.agent.approach(till_x=point[1], 272 | till_z=point[2], approach_dist=point[3], speed_threshold=velocity_threshold, 273 | rotation=angle) 274 | elif point[0] == "jump": 275 | if self.config_file["barrier_type"] in ["ramp", "platform"]: 276 | direction_x_z = np.array([-ramp_obj.target_position["x"] + self.agent.agent_position["x"], 277 | -ramp_obj.target_position["z"] + self.agent.agent_position["z"]]) 278 | direction_x_z = direction_x_z / np.linalg.norm(direction_x_z) 279 | 280 | self.agent.jump(height=point[1], width=point[2], direction_x_z=direction_x_z, 281 | hight_ground_height=ramp_obj.height) 282 | else: 283 | self.agent.jump(height=point[1], width=point[2]) 284 | elif point[0] == "settle": 285 | self.agent.settle_drop_object(point[1], stop_at_base=False) 286 | elif point[0] == "meta": 287 | meta_data = point 288 | # utils.make_video(self.agent.output_dir + "/images" + "_a", f"scene_a.mp4", os.path.join(self.agent.base_dir, f"path_{self.agent.path_no}")) 289 | # utils.make_video(self.agent.output_dir + "/images" + "_b", f"scene_b.mp4", os.path.join(self.agent.base_dir, f"path_{self.agent.path_no}")) 290 | # shutil.rmtree(self.agent.output_dir + "/images" + "_a") 291 | # shutil.rmtree(self.agent.output_dir + "/images" + "_b") 292 | if self.args.enable_image: 293 | utils.make_video(self.agent.output_dir + "_c", f"scene_c.mp4", 294 | os.path.join(self.agent.base_dir, f"path_{self.agent.path_no}")) 295 | shutil.rmtree(self.agent.output_dir + "_c") 296 | meta_data_arguments = { 297 | "config_file": self.config_file, 298 | "tdw_object": self.tdw_object, 299 | "high_scene": self.high_scene, 300 | "camera_obj": self.camera_obj, 301 | "target_1_color": self.target_1_color, 302 | "target_2_color": self.target_2_color, 303 | "target_1_id": self.target_1_id, 304 | "target_2_id": self.target_2_id, 305 | "obstacle_1": self.obstacle_1, 306 | "obstacle_2": self.obstacle_2, 307 | "ramp_1": self.ramp_1, 308 | "ramp_2": self.ramp_2, 309 | "target_1_position": self.target_1_position, 310 | "target_2_position": self.target_2_position, 311 | "agent_color": self.agent_color, 312 | "agent_shape": self.agent_shape, 313 | "agent": self.agent, 314 | "pit_ids": self.pit_ids, 315 | "pit_obj": self.pit_obj 316 | } 317 | object_meta = create_scene.construct_2_goal_object_meta_data(**meta_data_arguments) 318 | state_data = create_scene.create_state_data(self.agent, object_meta, self.camera_obj, meta_data) 319 | os.makedirs(os.path.join(self.agent.base_dir, f"path_{self.agent.path_no}"), exist_ok=True) 320 | with open(os.path.join(self.agent.base_dir, f"path_{self.agent.path_no}", "state_info.json"), "w") as fp: 321 | json.dump(state_data, fp) 322 | 323 | --------------------------------------------------------------------------------