├── .gitignore ├── LICENSE ├── README.md ├── composuite ├── __init__.py ├── arenas │ ├── compositional_arena.py │ ├── pick_place_arena.py │ ├── push_arena.py │ ├── shelf_arena.py │ └── trashcan_arena.py ├── assets │ ├── arenas │ │ ├── meshes │ │ │ ├── .DS_Store │ │ │ ├── trash_can.FCStd │ │ │ └── trash_can.stl │ │ ├── pick_place_arena.xml │ │ ├── push_arena.xml │ │ ├── shelf_arena.xml │ │ └── trashcan_arena.xml │ ├── objects │ │ ├── custombox-visual.xml │ │ ├── custombox.xml │ │ ├── dumbbell-visual.xml │ │ ├── dumbbell.xml │ │ ├── goal_wall.xml │ │ ├── hollowbox-visual.xml │ │ ├── hollowbox.xml │ │ ├── object_door_frame.xml │ │ ├── object_wall.xml │ │ ├── plate-visual.xml │ │ └── plate.xml │ └── textures │ │ ├── blue-wood.png │ │ ├── brass-ambra.png │ │ ├── bread.png │ │ ├── can.png │ │ ├── ceramic.png │ │ ├── cereal.png │ │ ├── clay.png │ │ ├── cream-plaster.png │ │ ├── dark-wood.png │ │ ├── dirt.png │ │ ├── glass.png │ │ ├── gray-felt.png │ │ ├── gray-plaster.png │ │ ├── gray-woodgrain.png │ │ ├── green-wood.png │ │ ├── lemon.png │ │ ├── light-wood.png │ │ ├── metal.png │ │ ├── pink-plaster.png │ │ ├── red-wood.png │ │ ├── steel-brushed.png │ │ ├── steel-scratched.png │ │ ├── white-bricks.png │ │ ├── white-plaster.png │ │ ├── wood-tiles.png │ │ ├── wood-varnished-panels.png │ │ └── yellow-plaster.png ├── demos │ ├── demo_compositional.py │ └── demo_compositional_teleop.py ├── env │ ├── compositional_env.py │ ├── gym_wrapper.py │ └── main.py ├── objects │ └── objects.py ├── tasks │ ├── pick_place_subtask.py │ ├── push_subtask.py │ ├── shelf_subtask.py │ └── trashcan_subtask.py └── utils │ ├── demo_utils.py │ └── mjk_utils.py ├── example_scripts └── train_ppo.py ├── requirements_default.txt ├── requirements_paper.txt ├── setup.py └── tests └── test_main_functions.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | .vscode/ 132 | 133 | #training folder 134 | spinningup_training/ 135 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Lifelong-ML 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # CompoSuite 2 | 3 | This repository contains the official release of 4 | [CompoSuite: A Compositional Reinforcement Learning Benchmark](https://arxiv.org/pdf/2207.04136.pdf). We include pre-trained models from our CoLLAs-22 paper in our accompanying data repository, [CompoSuite-Data](https://github.com/Lifelong-ML/CompoSuite-Data), and the multi-task and compositional PPO implementations in our adaptation of Spinning Up, [CompoSuite-spinningup](https://github.com/Lifelong-ML/CompoSuite-spinningup). 5 | 6 | CompoSuite is a benchmark of 256 compositional robotic manipulation tasks, each requiring fundamentally distinct behaviors. The diversity of tasks and their underlying compositional structure enables studying the ability of multi-task RL algorithms to extract the compositional properties of the environments, or simply explporing their ability to handle diverse tasks and generalize. 7 | 8 | If you use any part of CompoSuite for academic research, please cite our work using the following entry: 9 | 10 | ``` 11 | @inproceedings{mendez2022composuite 12 | title = {Compo{S}uite: A Compositional Reinforcement Learning Benchmark}, 13 | authors = {Jorge A. Mendez and Marcel Hussing and Meghna Gummadi and Eric Eaton}, 14 | booktitle = {1st Conference on Lifelong Learning Agents (CoLLAs-22)}, 15 | year = {2022} 16 | } 17 | ``` 18 | 19 | ## Installation 20 | 21 | We provide two sets of dependencies for installing CompoSuite: `requirements_default.txt` (for installing the latest compatible version of each dependency) and `requirements_paper.txt` (for reproducing the results from our CoLLAs-22 paper). To install, execute the following command sequence: 22 | 23 | CompoSuite requires Python<=3.11 (It might work with 3.12 but is not tested.). However, the example training scripts require Python version 3.6 or 3.7, for compatibility with Spinning Up. 24 | 25 | ``` 26 | git clone https://github.com/Lifelong-ML/CompoSuite.git 27 | cd CompoSuite 28 | pip install -r requirements_default.txt 29 | pip install -e . 30 | ``` 31 | 32 | To use the example training scripts, we provide a slightly modified version of Spinning Up which contains some minor changes to functionalities as well as the compositional learner, as described in our paper. This step is only required for running the example training scripts. To install our Spinning Up version, run: 33 | 34 | ``` 35 | git clone https://github.com/Lifelong-ML/CompoSuite-spinningup.git 36 | cd spinningup 37 | pip install -e . 38 | ``` 39 | 40 | 41 | ## Using CompoSuite 42 | 43 | To create an individual CompoSuite environment, execute: 44 | 45 | ``` 46 | composuite.make(ROBOT_NAME, OBJECT_NAME, OBSTACLE_NAME, OBJECTIVE_NAME, use_task_id_obs=False) 47 | ``` 48 | 49 | This returns a Gym environment that can be directly used for single-task training. The `use_task_id_obs` flag is used internally by CompoSuite to determine whether the observation should include a multi-hot indicator of the components that make up the task, as described in the paper. In the example above, we have set it to `False` because single-task agents make no use of the components. 50 | 51 | CompoSuite supports three problem settings in terms of how tasks are sampled, as described below: the full CompoSuite benchmark, small-scale CompoSuite∩<element> benchmarks, and restricted sampling CompoSuite\<element> benchmarks. Example code for how to train using Spinning Up under each of those settings, as well as a single-task setting, are included in `composuite/example_scripts/`. For multi-task training, make sure to set `use_task_id_obs=True`, unless you wish to withhold this information (note that with the default observation space, this is required to distinguish between tasks with a single observation). 52 | 53 | ### Full CompoSuite 54 | 55 | In this setting, tasks are sampled uniformly at random from all possible CompoSuite tasks. The following command creates two lists of tasks, one for training and one for testing zero-shot generalization: 56 | 57 | ``` 58 | train_tasks, test_tasks = composuite.sample_tasks(experiment_type='default', num_train=224) 59 | ``` 60 | 61 | Each task in the lists is a tuple `(ROBOT_NAME, OBJECT_NAME, OBSTACLE_NAME, OBJECTIVE_NAME)` that can be used to create a Gym environment with `composuite.make(...)`. 62 | 63 | ### Small-scale CompoSuite∩<element> 64 | 65 | This setting restricts the sampling to the set of tasks that contain the <element>. For example, CompoSuite∩IIWA considers only tasks that contain the IIWA arm. This restricts the sample to only 64 tasks, and limits the diversity of the tasks, both of which are useful during development. There are 16 possible choices of <element>, leading to 16 benchmarks of 64 tasks each. To sample training and evaluation tasks in this setting, execute: 66 | 67 | ``` 68 | train_tasks, test_tasks = composuite.sample_tasks(experiment_name='smallscale', smallscale_elem=, num_train=32) 69 | ``` 70 | 71 | ### Restricted CompoSuite\<element> 72 | 73 | In this last setting, the training set includes a single task with <element> and multiple uniformly sampled tasks without <element>, and the test set contains only the remaining tasks with <element>. This is the most challenging CompoSuite setting; for example, the agent might be required to solve all IIWA tasks after having trained on a single IIWA task. The following command creates the training and test set of tasks for this setting: 74 | 75 | ``` 76 | train_tasks, test_tasks = composuite.sample_tasks(experiment_name='holdout', holdout_elem=, num_train=56) 77 | ``` 78 | 79 | ## Using our example training scripts 80 | 81 | To train on a single environment using our example training scripts, run: 82 | 83 | ``` 84 | python -m example_scripts.train_ppo 85 | ``` 86 | 87 | To run something with multiple environments, the example scripts use MPI. For example, for training on 4 tasks with 4 MPI processes, run: 88 | 89 | ``` 90 | mpirun -np 4 python -m example_scripts.train_ppo --experiment-type smallscale --smallscale-elem IIWA --learner-type comp --num-train 4 91 | ``` 92 | 93 | Note that the multi-task and compositional PPO implementations in CompoSuite-spinningup assume that there is one MPI process per CompoSuite task (i.e., set `-np x` and `--num-train x` where `x` is the number of tasks). 94 | -------------------------------------------------------------------------------- /composuite/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | import robosuite 3 | from composuite.tasks.pick_place_subtask import PickPlaceSubtask 4 | from composuite.tasks.push_subtask import PushSubtask 5 | from composuite.tasks.shelf_subtask import ShelfSubtask 6 | from composuite.tasks.trashcan_subtask import TrashcanSubtask 7 | 8 | robosuite.environments.base.register_env(PickPlaceSubtask) 9 | robosuite.environments.base.register_env(PushSubtask) 10 | robosuite.environments.base.register_env(ShelfSubtask) 11 | robosuite.environments.base.register_env(TrashcanSubtask) 12 | 13 | from composuite.env.main import make, sample_tasks 14 | assets_root = os.path.join(os.path.dirname(__file__), "assets") 15 | -------------------------------------------------------------------------------- /composuite/arenas/compositional_arena.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from robosuite.models.arenas import Arena 4 | from robosuite.utils.mjcf_utils import array_to_string 5 | 6 | 7 | class CompositionalArena(Arena): 8 | """Workspace that contains two bins placed side by side. 9 | """ 10 | 11 | def __init__( 12 | self, fname, bin1_pos=None, bin2_pos=None 13 | ): 14 | """Initialize CompositionalArena 15 | 16 | Args: 17 | fname (str, Path): Path to the MJCF xml file. 18 | bin1_pos (list, optional): Position of the first bin in xyz 19 | coordinates. If None, position from default xml file is used. 20 | Defaults to None. 21 | bin2_pos (list, optional): Position of the first bin in xyz 22 | coordinates. If None, position from default xml file is used. 23 | Defaults to None. 24 | """ 25 | super().__init__(fname=fname) 26 | 27 | if bin1_pos is None: 28 | self.bin1_pos = np.array(list(map(float, self.worldbody.find( 29 | "./body[@name='bin1']").items()[1][1].split()))) 30 | else: 31 | self.bin1_pos = np.array(bin1_pos) 32 | self.worldbody.find( 33 | "./body[@name='bin1']").set("pos", array_to_string(bin1_pos)) 34 | 35 | if bin2_pos is None: 36 | self.bin2_pos = np.array(list(map(float, self.worldbody.find( 37 | "./body[@name='bin2']").items()[1][1].split()))) 38 | else: 39 | self.bin2_pos = np.array(bin2_pos) 40 | self.worldbody.find( 41 | "./body[@name='bin2']").set("pos", array_to_string(bin2_pos)) 42 | 43 | self.bin1_size = np.array(list(map(float, self.worldbody.find( 44 | "./body[@name='bin1']/geom[@name='floor1']").items()[2][1].split()))) 45 | self.bin2_size = np.array(list(map(float, self.worldbody.find( 46 | "./body[@name='bin2']/geom[@name='floor2']").items()[2][1].split()))) 47 | 48 | self.bin1_wall_right_size = np.array(list(map(float, self.worldbody.find( 49 | ".//geom[@name='bin1_right_wall']").items()[1][1].split()))) 50 | # self.bin1_wall_left_size = np.array(list(map(float, self.worldbody.find( 51 | # ".//geom[@name='bin1_left_wall']").items()[1][1].split()))) 52 | self.bin1_wall_front_size = np.array(list(map(float, self.worldbody.find( 53 | ".//geom[@name='bin1_front_wall']").items()[1][1].split()))) 54 | self.bin1_wall_rear_size = np.array(list(map(float, self.worldbody.find( 55 | ".//geom[@name='bin1_rear_wall']").items()[1][1].split()))) 56 | 57 | # TODO: do we need this line? 58 | self.table_top_abs = np.array(bin1_pos) 59 | 60 | self.configure_location() 61 | 62 | def configure_location(self): 63 | """Configures correct locations for this arena""" 64 | self.floor.set("pos", array_to_string(self.bottom_pos)) 65 | -------------------------------------------------------------------------------- /composuite/arenas/pick_place_arena.py: -------------------------------------------------------------------------------- 1 | from composuite.arenas.compositional_arena import CompositionalArena 2 | from composuite.utils.mjk_utils import xml_path_completion 3 | 4 | 5 | class PickPlaceArena(CompositionalArena): 6 | """Workspace for the PickPlace task. 7 | """ 8 | 9 | def __init__( 10 | self, bin1_pos=None, bin2_pos=None 11 | ): 12 | super().__init__(xml_path_completion("arenas/pick_place_arena.xml"), 13 | bin1_pos=bin1_pos, bin2_pos=bin2_pos) 14 | -------------------------------------------------------------------------------- /composuite/arenas/push_arena.py: -------------------------------------------------------------------------------- 1 | from composuite.utils.mjk_utils import xml_path_completion 2 | from composuite.arenas.compositional_arena import CompositionalArena 3 | 4 | 5 | class PushArena(CompositionalArena): 6 | """Workspace for the Push task. 7 | """ 8 | 9 | def __init__( 10 | self, bin1_pos=None, bin2_pos=None 11 | ): 12 | super().__init__(xml_path_completion("arenas/push_arena.xml"), bin1_pos=bin1_pos, 13 | bin2_pos=bin2_pos) 14 | 15 | 16 | -------------------------------------------------------------------------------- /composuite/arenas/shelf_arena.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from composuite.utils.mjk_utils import xml_path_completion 4 | from composuite.arenas.compositional_arena import CompositionalArena 5 | 6 | from robosuite.utils.mjcf_utils import array_to_string 7 | 8 | 9 | class ShelfArena(CompositionalArena): 10 | """Workspace for the Shelf task. 11 | """ 12 | 13 | 14 | def __init__( 15 | self, bin1_pos=None, bin2_pos=None 16 | ): 17 | super().__init__(xml_path_completion("arenas/shelf_arena.xml"), bin1_pos=bin1_pos, 18 | bin2_pos=bin2_pos) 19 | 20 | self.shelf_bottom_size = np.array(list(map(float, self.worldbody.find( 21 | ".//geom[@name='shelf_bottom_collision']").items()[1][1].split()))) 22 | self.shelf_top_size = np.array(list(map(float, self.worldbody.find( 23 | ".//geom[@name='shelf_top_collision']").items()[1][1].split()))) 24 | self.shelf_side_size = np.array(list(map(float, self.worldbody.find( 25 | ".//geom[@name='shelf_collision']").items()[1][1].split()))) 26 | 27 | self.shelf_bottom_pos = np.array(list(map(float, self.worldbody.find( 28 | ".//geom[@name='shelf_bottom_collision']").items()[0][1].split()))) 29 | self.shelf_top_pos = np.array(list(map(float, self.worldbody.find( 30 | ".//geom[@name='shelf_top_collision']").items()[0][1].split()))) 31 | self.shelf_side_pos = np.array(list(map(float, self.worldbody.find( 32 | ".//geom[@name='shelf_collision']").items()[0][1].split()))) 33 | 34 | self.shelf_pos = bin2_pos + np.array([0, self.shelf_top_size[1] + 0.05, 0]) 35 | self.worldbody.find( 36 | "./body[@name='shelf']").set("pos", array_to_string(self.shelf_pos)) 37 | 38 | 39 | self.shelf_x_size = self.shelf_top_size[0] 40 | self.shelf_y_size = self.shelf_top_size[1] 41 | self.shelf_thickness = self.shelf_top_size[2] 42 | self.shelf_z_offset = self.shelf_bottom_pos[2] + self.shelf_bottom_size[2] 43 | self.shelf_z_height = self.shelf_side_size[2] 44 | -------------------------------------------------------------------------------- /composuite/arenas/trashcan_arena.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from composuite.utils.mjk_utils import xml_path_completion 4 | from composuite.arenas.compositional_arena import CompositionalArena 5 | 6 | 7 | class TrashcanArena(CompositionalArena): 8 | """Workspace for the Trashcan task. 9 | """ 10 | 11 | def __init__( 12 | self, bin1_pos=None, bin2_pos=None 13 | ): 14 | super().__init__(xml_path_completion("arenas/trashcan_arena.xml"), bin1_pos=bin1_pos, 15 | bin2_pos=bin2_pos) 16 | 17 | self.trashcan_pos = np.array(list(map(float, self.worldbody.find( 18 | ".//body[@name='trashcan']").items()[1][1].split()))) 19 | self.trashcan_floor_size = np.array(list(map(float, self.worldbody.find( 20 | ".//geom[@name='trashcan5_collision']").items()[1][1].split()))) 21 | self.trashcan_height = np.array(list(map(float, self.worldbody.find( 22 | ".//geom[@name='trashcan1_collision']").items()[1][1].split())))[2] -------------------------------------------------------------------------------- /composuite/assets/arenas/meshes/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/arenas/meshes/.DS_Store -------------------------------------------------------------------------------- /composuite/assets/arenas/meshes/trash_can.FCStd: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/arenas/meshes/trash_can.FCStd -------------------------------------------------------------------------------- /composuite/assets/arenas/meshes/trash_can.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/arenas/meshes/trash_can.stl -------------------------------------------------------------------------------- /composuite/assets/arenas/pick_place_arena.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | -------------------------------------------------------------------------------- /composuite/assets/arenas/push_arena.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | -------------------------------------------------------------------------------- /composuite/assets/arenas/shelf_arena.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | 104 | 105 | 106 | 107 | -------------------------------------------------------------------------------- /composuite/assets/arenas/trashcan_arena.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 111 | -------------------------------------------------------------------------------- /composuite/assets/objects/custombox-visual.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | -------------------------------------------------------------------------------- /composuite/assets/objects/custombox.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | -------------------------------------------------------------------------------- /composuite/assets/objects/dumbbell-visual.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | -------------------------------------------------------------------------------- /composuite/assets/objects/dumbbell.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | -------------------------------------------------------------------------------- /composuite/assets/objects/goal_wall.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | -------------------------------------------------------------------------------- /composuite/assets/objects/hollowbox-visual.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | -------------------------------------------------------------------------------- /composuite/assets/objects/hollowbox.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | -------------------------------------------------------------------------------- /composuite/assets/objects/object_door_frame.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | -------------------------------------------------------------------------------- /composuite/assets/objects/object_wall.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | -------------------------------------------------------------------------------- /composuite/assets/objects/plate-visual.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | -------------------------------------------------------------------------------- /composuite/assets/objects/plate.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | -------------------------------------------------------------------------------- /composuite/assets/textures/blue-wood.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/blue-wood.png -------------------------------------------------------------------------------- /composuite/assets/textures/brass-ambra.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/brass-ambra.png -------------------------------------------------------------------------------- /composuite/assets/textures/bread.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/bread.png -------------------------------------------------------------------------------- /composuite/assets/textures/can.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/can.png -------------------------------------------------------------------------------- /composuite/assets/textures/ceramic.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/ceramic.png -------------------------------------------------------------------------------- /composuite/assets/textures/cereal.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/cereal.png -------------------------------------------------------------------------------- /composuite/assets/textures/clay.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/clay.png -------------------------------------------------------------------------------- /composuite/assets/textures/cream-plaster.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/cream-plaster.png -------------------------------------------------------------------------------- /composuite/assets/textures/dark-wood.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/dark-wood.png -------------------------------------------------------------------------------- /composuite/assets/textures/dirt.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/dirt.png -------------------------------------------------------------------------------- /composuite/assets/textures/glass.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/glass.png -------------------------------------------------------------------------------- /composuite/assets/textures/gray-felt.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/gray-felt.png -------------------------------------------------------------------------------- /composuite/assets/textures/gray-plaster.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/gray-plaster.png -------------------------------------------------------------------------------- /composuite/assets/textures/gray-woodgrain.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/gray-woodgrain.png -------------------------------------------------------------------------------- /composuite/assets/textures/green-wood.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/green-wood.png -------------------------------------------------------------------------------- /composuite/assets/textures/lemon.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/lemon.png -------------------------------------------------------------------------------- /composuite/assets/textures/light-wood.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/light-wood.png -------------------------------------------------------------------------------- /composuite/assets/textures/metal.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/metal.png -------------------------------------------------------------------------------- /composuite/assets/textures/pink-plaster.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/pink-plaster.png -------------------------------------------------------------------------------- /composuite/assets/textures/red-wood.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/red-wood.png -------------------------------------------------------------------------------- /composuite/assets/textures/steel-brushed.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/steel-brushed.png -------------------------------------------------------------------------------- /composuite/assets/textures/steel-scratched.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/steel-scratched.png -------------------------------------------------------------------------------- /composuite/assets/textures/white-bricks.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/white-bricks.png -------------------------------------------------------------------------------- /composuite/assets/textures/white-plaster.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/white-plaster.png -------------------------------------------------------------------------------- /composuite/assets/textures/wood-tiles.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/wood-tiles.png -------------------------------------------------------------------------------- /composuite/assets/textures/wood-varnished-panels.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/wood-varnished-panels.png -------------------------------------------------------------------------------- /composuite/assets/textures/yellow-plaster.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Lifelong-ML/CompoSuite/1fa36f67f31aeccc9ef75748bfc797960e044a86/composuite/assets/textures/yellow-plaster.png -------------------------------------------------------------------------------- /composuite/demos/demo_compositional.py: -------------------------------------------------------------------------------- 1 | import composuite 2 | from composuite.utils.demo_utils import * 3 | from robosuite.utils.input_utils import * 4 | 5 | 6 | def simulate_compositional_env(): 7 | """Initialize a compositional environment and run with 8 | random actions for visualization. 9 | """ 10 | robot = choose_robots(exclude_bimanual=True) 11 | task = choose_task() 12 | obj = choose_object() 13 | obstacle = choose_obstacle() 14 | if obstacle == "None": obstacle = None 15 | 16 | # Help message to user 17 | print() 18 | print("Press \"H\" to show the viewer control panel.") 19 | 20 | # initialize the task 21 | env = composuite.make( 22 | robot=robot, 23 | obj=obj, 24 | obstacle=obstacle, 25 | task=task, 26 | has_renderer=True, 27 | ignore_done=True, 28 | ) 29 | env.reset() 30 | env.viewer.set_camera(camera_id=1) 31 | 32 | # Get action limits 33 | low, high = env.action_spec 34 | 35 | # do visualization 36 | for _ in range(10000): 37 | action = np.random.uniform(low, high) 38 | env.step(action) 39 | env.render() 40 | 41 | if __name__ == "__main__": 42 | simulate_compositional_env() 43 | -------------------------------------------------------------------------------- /composuite/demos/demo_compositional_teleop.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | 4 | import composuite 5 | from composuite.utils.demo_utils import * 6 | 7 | from robosuite.utils.input_utils import input2action 8 | from robosuite.wrappers import VisualizationWrapper 9 | 10 | 11 | def parse_demo_args(): 12 | """Parse input arguments for teleop script. 13 | 14 | Returns: 15 | Namespace: Parsed arguments. 16 | """ 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument("--robots", nargs="+", type=str, default="Panda", help="Which robot(s) to use in the env") 19 | parser.add_argument("--config", type=str, default="single-arm-opposed", 20 | help="Specified environment configuration if necessary") 21 | parser.add_argument("--arm", type=str, default="right", help="Which arm to control (eg bimanual) 'right' or 'left'") 22 | parser.add_argument("--switch-on-grasp", action="store_true", help="Switch gripper control on gripper action") 23 | parser.add_argument("--toggle-camera-on-grasp", action="store_true", help="Switch camera angle on gripper action") 24 | parser.add_argument("--controller", type=str, default="osc", help="Choice of controller. Can be 'ik' or 'osc'") 25 | parser.add_argument("--device", type=str, default="keyboard") 26 | parser.add_argument("--pos-sensitivity", type=float, default=1.0, help="How much to scale position user inputs") 27 | parser.add_argument("--rot-sensitivity", type=float, default=1.0, help="How much to scale rotation user inputs") 28 | parser.add_argument("--shelf-reward", default='default') 29 | args = parser.parse_args() 30 | 31 | return args 32 | 33 | 34 | def simulate_compositional_env_teleop(): 35 | """Run code to manually teleoperate a chosen environment. 36 | """ 37 | args = parse_demo_args() 38 | 39 | robot = choose_robots(exclude_bimanual=True) 40 | task = choose_task() 41 | obj = choose_object() 42 | obstacle = choose_obstacle() 43 | if obstacle == "None": obstacle = None 44 | 45 | # Create environment 46 | env = composuite.make( 47 | robot=robot, 48 | obj=obj, 49 | obstacle=obstacle, 50 | task=task, 51 | has_renderer=True, 52 | # render_camera="agentview", 53 | ignore_done=True, 54 | controller="osc_pose", 55 | hard_reset=False, 56 | ) 57 | 58 | # Wrap this environment in a visualization wrapper 59 | env = VisualizationWrapper(env, indicator_configs=None) 60 | 61 | # Setup printing options for numbers 62 | np.set_printoptions(formatter={'float': lambda x: "{0:0.3f}".format(x)}) 63 | 64 | # initialize device 65 | if args.device == "keyboard": 66 | from robosuite.devices import Keyboard 67 | 68 | device = Keyboard(pos_sensitivity=args.pos_sensitivity, rot_sensitivity=args.rot_sensitivity) 69 | env.viewer.add_keypress_callback("any", device.on_press) 70 | env.viewer.add_keyup_callback("any", device.on_release) 71 | env.viewer.add_keyrepeat_callback("any", device.on_press) 72 | elif args.device == "spacemouse": 73 | from robosuite.devices import SpaceMouse 74 | 75 | device = SpaceMouse(pos_sensitivity=args.pos_sensitivity, rot_sensitivity=args.rot_sensitivity) 76 | else: 77 | raise Exception( 78 | "Invalid device choice: choose either 'keyboard' or 'spacemouse'." 79 | ) 80 | 81 | while True: 82 | # Reset the environment 83 | obs = env.reset() 84 | 85 | # Setup rendering 86 | cam_id = 0 87 | num_cam = len(env.sim.model.camera_names) 88 | env.render() 89 | 90 | # Initialize variables that should the maintained between resets 91 | last_grasp = 0 92 | 93 | # Initialize device control 94 | device.start_control() 95 | 96 | while True: 97 | # Set active robot 98 | active_robot = env.robots[0] if args.config == "bimanual" else env.robots[args.arm == "left"] 99 | 100 | # Get the newest action 101 | action, grasp = input2action( 102 | device=device, 103 | robot=active_robot, 104 | active_arm=args.arm, 105 | env_configuration=args.config 106 | ) 107 | 108 | # If action is none, then this a reset so we should break 109 | if action is None: 110 | break 111 | 112 | # If the current grasp is active (1) and last grasp is not (-1) (i.e.: grasping input just pressed), 113 | # toggle arm control and / or camera viewing angle if requested 114 | if last_grasp < 0 < grasp: 115 | if args.switch_on_grasp: 116 | args.arm = "left" if args.arm == "right" else "right" 117 | if args.toggle_camera_on_grasp: 118 | cam_id = (cam_id + 1) % num_cam 119 | env.viewer.set_camera(camera_id=cam_id) 120 | # Update last grasp 121 | last_grasp = grasp 122 | 123 | # Fill out the rest of the action space if necessary 124 | rem_action_dim = env.action_dim - action.size 125 | if rem_action_dim > 0: 126 | # Initialize remaining action space 127 | rem_action = np.zeros(rem_action_dim) 128 | # This is a multi-arm setting, choose which arm to control and fill the rest with zeros 129 | if args.arm == "right": 130 | action = np.concatenate([action, rem_action]) 131 | elif args.arm == "left": 132 | action = np.concatenate([rem_action, action]) 133 | else: 134 | # Only right and left arms supported 135 | print("Error: Unsupported arm specified -- " 136 | "must be either 'right' or 'left'! Got: {}".format(args.arm)) 137 | elif rem_action_dim < 0: 138 | # We're in an environment with no gripper action space, so trim the action space to be the action dim 139 | action = action[:env.action_dim] 140 | 141 | # Step through the simulation and render 142 | _, reward, _, _ = env.step(action) 143 | print('Reward: {}'.format(reward)) 144 | env.render() 145 | 146 | 147 | 148 | if __name__ == "__main__": 149 | simulate_compositional_env_teleop() -------------------------------------------------------------------------------- /composuite/env/compositional_env.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from composuite.objects.objects import GoalWallObject, ObjectWallObject, ObjectDoorFrameObject 4 | from composuite.objects.objects import \ 5 | DumbbellObject, DumbbellVisualObject, \ 6 | PlateObject, PlateVisualObject, \ 7 | HollowBoxObject, HollowBoxVisualObject, \ 8 | CustomBoxObject, CustomBoxVisualObject 9 | 10 | import robosuite.utils.transform_utils as T 11 | from robosuite.environments.manipulation.single_arm_env import SingleArmEnv 12 | from robosuite.utils.observables import Observable, sensor 13 | from robosuite.utils.transform_utils import convert_quat 14 | from robosuite.utils.placement_samplers import SequentialCompositeSampler, UniformRandomSampler 15 | 16 | from robosuite.models.tasks import ManipulationTask 17 | 18 | from robosuite.models.grippers import GripperModel 19 | from robosuite.models.base import MujocoModel 20 | 21 | from collections import OrderedDict 22 | 23 | 24 | class CompositionalEnv(SingleArmEnv): 25 | def __init__( 26 | self, 27 | robots, # first compositional axis 28 | object_type, 29 | obstacle, 30 | bin1_pos, 31 | bin2_pos, 32 | env_configuration="default", 33 | controller_configs=None, 34 | mount_types="default", 35 | gripper_types="RethinkGripper", 36 | initialization_noise=None, 37 | use_camera_obs=True, 38 | use_object_obs=True, 39 | use_goal_obs=True, 40 | use_obstacle_obs=True, 41 | use_task_id_obs=False, 42 | has_renderer=False, 43 | has_offscreen_renderer=True, 44 | render_camera="frontview", 45 | render_collision_mesh=False, 46 | render_visual_mesh=True, 47 | render_gpu_device_id=-1, 48 | control_freq=20, 49 | horizon=1000, 50 | ignore_done=True, 51 | hard_reset=True, 52 | camera_names="agentview", 53 | camera_heights=256, 54 | camera_widths=256, 55 | camera_depths=False, 56 | reward_scale=1.0, 57 | reward_shaping=True, 58 | ): 59 | """Initializes a base compositional environment. Most of the input 60 | variables are passed through to robosuite directly. Variables that 61 | have to be specified specifically for the compositional environment 62 | are the following. 63 | 64 | Args: 65 | obstacle (str): Obstacle to use in the environment. 66 | bin1_pos (list, np.array): Position of the first bin in xyz. 67 | bin2_pos (list, np.array): Position of the second bin in xyz. 68 | use_object_obs (bool, optional): When True, return object 69 | position in observation. Defaults to True. 70 | use_goal_obs (bool, optional): When True, return goal position 71 | in observation. Defaults to True. 72 | use_obstacle_obs (bool, optional):When True, return obstacle 73 | position in observation. Defaults to True. 74 | use_task_id_obs (bool, optional): When True, include task 75 | indicator as one-hot vector in observation. Defaults to False. 76 | reward_scale (float, optional): Multiplier for default rewards. 77 | Defaults to 1.0. 78 | reward_shaping (bool, optional): When True, use reward shaping 79 | instead of sparse rewards. Defaults to True. 80 | 81 | Raises: 82 | ValueError: Raises if you try to pass in a different gripper. 83 | This option is left in for potential future development purposes. 84 | """ 85 | if gripper_types != "RethinkGripper": 86 | raise ValueError( 87 | 'Cannot use a different gripper to ensure that observation \ 88 | spaces are the same. Please use RethinkgGripper.') 89 | self.step_counter = 0 90 | 91 | # task settings 92 | self.object_type = object_type 93 | self.obstacle_type = None if obstacle == "None" else obstacle 94 | self.robot_name = robots 95 | self.obj_class_map = OrderedDict({ 96 | "Box": (CustomBoxObject, CustomBoxVisualObject), 97 | "Dumbbell": (DumbbellObject, DumbbellVisualObject), 98 | "Plate": (PlateObject, PlateVisualObject), 99 | "Hollowbox": (HollowBoxObject, HollowBoxVisualObject)}) 100 | self.obstacle_class_map = OrderedDict({ 101 | None: None, 102 | "ObjectWall": ObjectWallObject, 103 | "ObjectDoor": ObjectDoorFrameObject, 104 | "GoalWall": GoalWallObject}) 105 | self.robot_list = ["IIWA", "Jaco", "Panda", "Kinova3"] 106 | 107 | # Needed for Robosuite gym wrapper 108 | self.use_object_obs = use_object_obs 109 | self.use_goal_obs = use_goal_obs 110 | self.use_obstacle_obs = use_obstacle_obs 111 | self.use_task_id_obs = use_task_id_obs 112 | 113 | # whether to use ground-truth object states 114 | self.object_sensor_names = None 115 | 116 | # reward configuration 117 | self.reward_scale = reward_scale 118 | self.reward_shaping = reward_shaping 119 | 120 | # settings for bin position 121 | self.bin1_pos = bin1_pos 122 | self.bin2_pos = bin2_pos 123 | 124 | super().__init__( 125 | robots, 126 | env_configuration=env_configuration, 127 | controller_configs=controller_configs, 128 | mount_types=mount_types, 129 | gripper_types=gripper_types, 130 | initialization_noise=initialization_noise, 131 | use_camera_obs=use_camera_obs, 132 | has_renderer=has_renderer, 133 | has_offscreen_renderer=has_offscreen_renderer, 134 | render_camera=render_camera, 135 | render_collision_mesh=render_collision_mesh, 136 | render_visual_mesh=render_visual_mesh, 137 | render_gpu_device_id=render_gpu_device_id, 138 | control_freq=control_freq, 139 | horizon=horizon, 140 | ignore_done=ignore_done, 141 | hard_reset=hard_reset, 142 | camera_names=camera_names, 143 | camera_heights=camera_heights, 144 | camera_widths=camera_widths, 145 | camera_depths=camera_depths, 146 | ) 147 | 148 | def staged_reward(self, action=None): 149 | """Dense reward function that depends on the task. Needs to be 150 | implemented on task level. 151 | """ 152 | raise NotImplementedError( 153 | "Needs to be implemented in inheriting class.") 154 | 155 | def _load_model(self): 156 | """ 157 | Loads an xml model, puts it in self.model 158 | """ 159 | super()._load_model() 160 | 161 | def _initialize_model(self): 162 | """Load all the required arena model information and store it 163 | in class variables. Create the ManipulationTask that will hold 164 | all the information about arena, robot, and objects of interest. 165 | """ 166 | self.bin1_pos = self.mujoco_arena.bin1_pos 167 | self.bin2_pos = self.mujoco_arena.bin2_pos 168 | 169 | self.bin1_size = self.mujoco_arena.bin1_size 170 | self.bin2_size = self.mujoco_arena.bin2_size 171 | 172 | self.mujoco_arena.set_origin([0, 0, 0]) 173 | 174 | # Adjust base pose accordingly 175 | xpos = self.robots[0].robot_model.base_xpos_offset["bins"] 176 | self.robots[0].robot_model.set_base_xpos(xpos) 177 | 178 | self.object = self.obj_class_map[self.object_type][0]( 179 | self.object_type.title()) 180 | self.visual_object = self.obj_class_map[self.object_type][1]( 181 | 'Visual' + self.object_type.title()) 182 | mujoco_objects = [self.visual_object, self.object] 183 | 184 | if self.obstacle_type is not None: 185 | self.obstacle = self.obstacle_class_map[self.obstacle_type]( 186 | 'Obstacle' + self.obstacle_type) 187 | mujoco_objects.append(self.obstacle) 188 | 189 | self.model = ManipulationTask( 190 | mujoco_arena=self.mujoco_arena, 191 | mujoco_robots=[ 192 | robot.robot_model for robot in self.robots], 193 | mujoco_objects=mujoco_objects, 194 | ) 195 | 196 | def _setup_references(self): 197 | """ 198 | Sets up references to important components. A reference is typically an 199 | index or a list of indices that point to the corresponding elements 200 | in a flatten array, which is how MuJoCo stores physical simulation data. 201 | """ 202 | super()._setup_references() 203 | 204 | # object-specific ids 205 | self.obj_body_id = self.sim.model.body_name2id( 206 | self.object.root_body) 207 | self.visual_obj_body_id = self.sim.model.body_name2id( 208 | self.visual_object.root_body) 209 | # required for setting up observations in parent task 210 | self.goal_body_id = self.visual_obj_body_id 211 | 212 | if self.obstacle_type is not None: 213 | self.obstacle_body_id = self.sim.model.body_name2id( 214 | self.obstacle.root_body) 215 | 216 | def _reset_internal(self): 217 | """Resets simulation internal configurations. 218 | """ 219 | super()._reset_internal() 220 | # Reset all object positions using initializer sampler if we're not directly loading from an xml 221 | if not self.deterministic_reset: 222 | 223 | # Sample from the placement initializer for all objects 224 | object_placements = self.placement_initializer.sample() 225 | 226 | # Loop through all objects and reset their positions 227 | for obj_pos, obj_quat, obj in object_placements.values(): 228 | # Set the visual object body locations 229 | if "visual" in obj.name.lower(): 230 | self.sim.model.body_pos[self.visual_obj_body_id] = obj_pos 231 | self.sim.model.body_quat[self.visual_obj_body_id] = obj_quat 232 | elif "obstacle" in obj.name.lower(): 233 | # Set the obstacle body locations 234 | self.sim.model.body_pos[self.obstacle_body_id] = obj_pos 235 | self.sim.model.body_quat[self.obstacle_body_id] = obj_quat 236 | else: 237 | # Set the collision object joints 238 | self.sim.data.set_joint_qpos( 239 | obj.joints[0], np.concatenate( 240 | [np.array(obj_pos), np.array(obj_quat)])) 241 | 242 | # Make sure to update sensors' active and enabled states 243 | if self.object_sensor_names is not None: 244 | for name in self.object_sensor_names: 245 | # Set all of these sensors to be enabled and active if this is the active object, else False 246 | self._observables[name].set_enabled(True) 247 | self._observables[name].set_active(True) 248 | 249 | def _setup_observables(self): 250 | """Create all existing robosuite observables. 251 | Extrends information from upper level classes. 252 | Here, specifically add task, object, goal and obstacle 253 | observables for the compositional task by looping through 254 | all components, creating the Observables and extending the 255 | Observables dict. 256 | 257 | Returns: 258 | dict: 259 | """ 260 | observables = super()._setup_observables() 261 | # return self.task._setup_observables(observables) 262 | pf = self.robots[0].robot_model.naming_prefix 263 | 264 | # for conversion to relative gripper frame 265 | @sensor(modality="ref") 266 | def world_pose_in_gripper(obs_cache): 267 | return T.pose_inv(T.pose2mat((obs_cache[f"{pf}eef_pos"], obs_cache[f"{pf}eef_quat"]))) if\ 268 | f"{pf}eef_pos" in obs_cache and f"{pf}eef_quat" in obs_cache else np.eye(4) 269 | sensors = [world_pose_in_gripper] 270 | names = ["world_pose_in_gripper"] 271 | enableds = [True] 272 | actives = [False] 273 | 274 | if self.use_task_id_obs: 275 | modality = "task" 276 | 277 | @sensor(modality=modality) 278 | def subtask_id(obs_cache): 279 | onehot = np.zeros(4) 280 | onehot[self.subtask_id] = 1 281 | return onehot 282 | 283 | @sensor(modality=modality) 284 | def object_id(obs_cache): 285 | onehot = np.zeros(len(self.obj_class_map)) 286 | obj_id = list(self.obj_class_map.keys() 287 | ).index(self.object_type) 288 | onehot[obj_id] = 1 289 | return onehot 290 | 291 | @sensor(modality=modality) 292 | def obstacle_id(obs_cache): 293 | onehot = np.zeros(len(self.obstacle_class_map)) 294 | obstacle_id = list(self.obstacle_class_map.keys()).index( 295 | self.obstacle_type) 296 | onehot[obstacle_id] = 1 297 | return onehot 298 | 299 | @sensor(modality=modality) 300 | def robot_id(obs_cache): 301 | onehot = np.zeros(len(self.robot_list)) 302 | robot_id = self.robot_list.index(self.robot_name) 303 | onehot[robot_id] = 1 304 | return onehot 305 | 306 | sensors += [subtask_id, object_id, obstacle_id, robot_id] 307 | names += ['subtask_id', 'object_id', 308 | 'obstacle_id', 'robot_id'] 309 | 310 | enableds += [True] * 4 311 | actives += [True] * 4 312 | 313 | if self.use_object_obs: 314 | # Get robot prefix and define observables modality 315 | modality = "object" 316 | 317 | # object-related observables 318 | @sensor(modality=modality) 319 | def obj_pos(obs_cache): 320 | return np.array(self.sim.data.body_xpos[self.obj_body_id]) 321 | 322 | @sensor(modality=modality) 323 | def obj_quat(obs_cache): 324 | return convert_quat(np.array(self.sim.data.body_xquat[self.obj_body_id]), to="xyzw") 325 | 326 | @sensor(modality=modality) 327 | def obj_to_eef_pos(obs_cache): 328 | # Immediately return default value if cache is empty 329 | if any([name not in obs_cache for name in 330 | ["obj_pos", "obj_quat", "world_pose_in_gripper"]]): 331 | return np.zeros(3) 332 | obj_pose = T.pose2mat( 333 | (obs_cache["obj_pos"], obs_cache["obj_quat"])) 334 | rel_pose = T.pose_in_A_to_pose_in_B( 335 | obj_pose, obs_cache["world_pose_in_gripper"]) 336 | rel_pos, rel_quat = T.mat2pose(rel_pose) 337 | obs_cache[f"obj_to_{pf}eef_quat"] = rel_quat 338 | return rel_pos 339 | 340 | @sensor(modality=modality) 341 | def obj_to_eef_quat(obs_cache): 342 | return obs_cache[f"obj_to_{pf}eef_quat"] if \ 343 | f"obj_to_{pf}eef_quat" in obs_cache else np.zeros(4) 344 | 345 | sensors += [obj_pos, obj_quat, obj_to_eef_pos, obj_to_eef_quat] 346 | names += ['obj_pos', 'obj_quat', 347 | 'obj_to_eef_pos', 'obj_to_eef_quat'] 348 | 349 | enableds += [True] * 4 350 | actives += [True] * 4 351 | 352 | if self.use_goal_obs: 353 | modality = "goal" 354 | 355 | # goal-related observables 356 | @sensor(modality=modality) 357 | def goal_pos(obs_cache): 358 | return np.array(self.sim.data.body_xpos[self.goal_body_id]) 359 | 360 | @sensor(modality=modality) 361 | def goal_quat(obs_cache): 362 | return convert_quat(np.array(self.sim.data.body_xquat[self.goal_body_id]), to="xyzw") 363 | 364 | @sensor(modality=modality) 365 | def goal_to_eef_pos(obs_cache): 366 | # Immediately return default value if cache is empty 367 | if any([name not in obs_cache for name in 368 | ["goal_pos", "goal_quat", "world_pose_in_gripper"]]): 369 | return np.zeros(3) 370 | goal_pose = T.pose2mat( 371 | (obs_cache["goal_pos"], obs_cache["goal_quat"])) 372 | rel_pose = T.pose_in_A_to_pose_in_B( 373 | goal_pose, obs_cache["world_pose_in_gripper"]) 374 | rel_pos, rel_quat = T.mat2pose(rel_pose) 375 | obs_cache[f"goal_to_{pf}eef_quat"] = rel_quat 376 | return rel_pos 377 | 378 | @sensor(modality=modality) 379 | def goal_to_eef_quat(obs_cache): 380 | return obs_cache[f"goal_to_{pf}eef_quat"] if \ 381 | f"goal_to_{pf}eef_quat" in obs_cache else np.zeros(4) 382 | 383 | # in principle, like other things, we could add the quat as well 384 | @sensor(modality=modality) 385 | def obj_to_goal(obs_cache): 386 | return obs_cache["goal_pos"] - obs_cache["obj_pos"] if \ 387 | "obj_pos" in obs_cache and "goal_pos" in obs_cache else np.zeros(3) 388 | 389 | sensors += [goal_pos, goal_quat, goal_to_eef_pos, 390 | goal_to_eef_quat, obj_to_goal] 391 | names += ['goal_pos', 'goal_quat', 'goal_to_eef_pos', 392 | 'goal_to_eef_quat', 'obj_to_goal'] 393 | 394 | enableds += [True] * 5 395 | actives += [True] * 5 396 | 397 | if self.use_obstacle_obs: 398 | modality = "obstacle" 399 | 400 | # goal-related observables 401 | @sensor(modality=modality) 402 | def obstacle_pos(obs_cache): 403 | return np.array(self.sim.data.body_xpos[self.obstacle_body_id]) \ 404 | if self.obstacle_type is not None else np.zeros(3) 405 | 406 | @sensor(modality=modality) 407 | def obstacle_quat(obs_cache): 408 | return convert_quat(np.array(self.sim.data.body_xquat[self.obstacle_body_id]), to="xyzw") \ 409 | if self.obstacle_type is not None else np.zeros(4) 410 | 411 | @sensor(modality=modality) 412 | def obstacle_to_eef_pos(obs_cache): 413 | # Immediately return default value if cache is empty 414 | if any([name not in obs_cache for name in 415 | ["obstacle_pos", "obstacle_quat", "world_pose_in_gripper"]]): 416 | return np.zeros(3) 417 | obstacle_pose = T.pose2mat( 418 | (obs_cache["obstacle_pos"], obs_cache["obstacle_quat"])) 419 | rel_pose = T.pose_in_A_to_pose_in_B( 420 | obstacle_pose, obs_cache["world_pose_in_gripper"]) 421 | rel_pos, rel_quat = T.mat2pose(rel_pose) 422 | obs_cache[f"obstacle_to_{pf}eef_quat"] = rel_quat 423 | return rel_pos 424 | 425 | @sensor(modality=modality) 426 | def obstacle_to_eef_quat(obs_cache): 427 | return obs_cache[f"obstacle_to_{pf}eef_quat"] if \ 428 | f"obstacle_to_{pf}eef_quat" in obs_cache else np.zeros(4) 429 | 430 | sensors += [obstacle_pos, obstacle_quat, 431 | obstacle_to_eef_pos, obstacle_to_eef_quat] 432 | names += ['obstacle_pos', 'obstacle_quat', 433 | 'obstacle_to_eef_pos', 'obstacle_to_eef_quat'] 434 | 435 | enableds += [True] * 4 436 | actives += [True] * 4 437 | 438 | # Create observables 439 | for name, s, enabled, active in zip(names, sensors, enableds, actives): 440 | observables[name] = Observable( 441 | name=name, 442 | sensor=s, 443 | sampling_rate=self.control_freq, 444 | enabled=enabled, 445 | active=active 446 | ) 447 | return observables 448 | 449 | def _check_success(self): 450 | """Check whether task was solved. This is task specific and needs to 451 | be implemented on task level. 452 | """ 453 | raise NotImplementedError('Has to be implemented in inheriting class.') 454 | 455 | def _get_placement_initializer(self): 456 | """Helper function for defining placement initializer 457 | and object sampling bounds. 458 | """ 459 | self.placement_initializer = SequentialCompositeSampler( 460 | name="ObjectSampler") 461 | 462 | # Add obstacle 463 | if self.obstacle_type is not None: 464 | if 'object' in self.obstacle_type.lower(): 465 | # Place obstacle a quarter of the way into the table, leaving ~3/4 of the table to place the object 466 | obstacle_pos = [-self.bin1_size[0] / 467 | 2.2, 0, 0] + self.bin1_pos 468 | rotation = np.pi / 2 469 | 470 | # sample anywhere beyond the obstacle 471 | obj_x_min = 0 472 | obj_x_max = self.bin1_size[0] / 2 473 | obj_y_min = -(self.bin1_size[1] / 2) 474 | obj_y_max = -obj_y_min 475 | 476 | elif 'goal' in self.obstacle_type.lower(): 477 | obstacle_pos = [ 478 | 0, self.bin1_size[1], 0] + self.bin1_pos 479 | rotation = 0. 480 | obj_x_min = -(self.bin1_size[0] / 2) 481 | obj_x_max = -obj_x_min 482 | obj_y_min = -(self.bin1_size[1] / 2) 483 | obj_y_max = -obj_y_min 484 | else: 485 | raise ValueError( 486 | 'Obstacle must be placed between robot and object for object obstacles.') 487 | 488 | self.placement_initializer.append_sampler( 489 | sampler=UniformRandomSampler( 490 | name="CollisionObstacleSampler", 491 | mujoco_objects=self.obstacle, 492 | rotation=rotation, 493 | rotation_axis='z', 494 | ensure_object_boundary_in_range=False, 495 | ensure_valid_placement=True, 496 | reference_pos=obstacle_pos 497 | ) 498 | ) 499 | else: 500 | # can sample anywhere in bin 501 | obj_x_min = -(self.bin1_size[0] / 2) 502 | obj_x_max = -obj_x_min 503 | obj_y_min = -(self.bin1_size[1] / 2) 504 | obj_y_max = -obj_y_min 505 | 506 | # each object should just be sampled in the bounds of the shelf (with some tolerance) 507 | self.placement_initializer.append_sampler( 508 | sampler=UniformRandomSampler( 509 | name="CollisionObjectSampler", 510 | mujoco_objects=self.object, 511 | x_range=[obj_x_min, obj_x_max], 512 | y_range=[obj_y_min, obj_y_max], 513 | rotation=0, 514 | rotation_axis='z', 515 | ensure_object_boundary_in_range=True, 516 | ensure_valid_placement=True, 517 | reference_pos=self.bin1_pos, 518 | z_offset=0.02, 519 | ) 520 | ) 521 | 522 | def reward(self, action=None): 523 | """Calls the specific reward function for the task. 524 | Sparse un-normalized reward: 525 | - a discrete reward of 1.0 if the object is placed in the shelf. 526 | Use Un-normalized components if using reward shaping, 527 | where the maximum is returned if not solved 528 | - stages are defined in the specific task. 529 | 530 | Note that a successfully completed task (object in shelf) will 531 | return 1.0 irregardless of whether the environment is using 532 | sparse or shaped rewards. 533 | 534 | Args: 535 | action (np.array): Action to be passed to the staged rewards. This may 536 | be required in some cases to compute the reward. 537 | 538 | Returns: 539 | float: reward value 540 | """ 541 | # compute sparse rewards 542 | reward = float(self._check_success()) 543 | # add in shaped re wards 544 | if reward < 1 and self.reward_shaping: 545 | staged_rewards = self.staged_rewards(action) 546 | reward += max(staged_rewards) 547 | self.step_counter += 1 548 | if self.reward_scale is not None: 549 | reward *= self.reward_scale 550 | return reward 551 | 552 | def visualize(self, vis_settings): 553 | super().visualize(vis_settings=vis_settings) 554 | # Color the gripper visualization site according to its distance to the closest object 555 | if vis_settings["grippers"]: 556 | # Visualize the distance to this target 557 | self._visualize_gripper_to_target( 558 | gripper=self.robots[0].gripper, 559 | target=self.object.root_body, 560 | target_type="body", 561 | ) 562 | 563 | def _check_grasp(self, gripper, object_geoms): 564 | """Check whether gripper is grasping any object other than 565 | the plate. 566 | 567 | Args: 568 | gripper (GripperModel or str or list of str or list of list of str): 569 | Mujoco gripper model as specified by robosuite. 570 | object_geoms (str or list of str or MujocoModel): Object geoms to check 571 | for collsions. 572 | 573 | Returns: 574 | bool: True if the gripper is grasping the object 575 | """ 576 | if self.object_type.title() == "Plate": 577 | return self._check_plate_grasp(gripper, object_geoms) 578 | return super()._check_grasp(gripper, object_geoms) 579 | 580 | def _check_plate_grasp(self, gripper, object_geoms): 581 | """ 582 | Checks whether the gripper is grasping the plate in the environement. 583 | This will return True if at one finger and the opposing sides fingerpad 584 | or both fingerpads are in contact with the plate. 585 | 586 | Args: 587 | gripper (GripperModel or str or list of str or list of list of str): 588 | Mujoco gripper model as specified by robosuite. 589 | object_geoms (str or list of str or MujocoModel): Object geoms to check 590 | for collsions. 591 | 592 | Returns: 593 | bool: True if the gripper is grasping the given object. 594 | """ 595 | # Convert object, gripper geoms into standardized form 596 | if isinstance(object_geoms, MujocoModel): 597 | o_geoms = object_geoms.contact_geoms 598 | else: 599 | o_geoms = [object_geoms] if type( 600 | object_geoms) is str else object_geoms 601 | if isinstance(gripper, GripperModel): 602 | right_g_geoms = [gripper.important_geoms["right_finger"], 603 | gripper.important_geoms["left_fingerpad"]] 604 | left_g_geoms = [gripper.important_geoms["left_finger"], 605 | gripper.important_geoms["right_fingerpad"]] 606 | elif type(gripper) is str: 607 | right_g_geoms = [[gripper]] 608 | left_g_geoms = [[gripper]] 609 | else: 610 | # Parse each element in the gripper_geoms list accordingly 611 | right_g_geoms = [[g_group] if type( 612 | g_group) is str else g_group for g_group in gripper] 613 | left_g_geoms = [[g_group] if type( 614 | g_group) is str else g_group for g_group in gripper] 615 | 616 | # Search for collisions between each gripper geom group and the object geoms group 617 | for right_g_group in right_g_geoms: 618 | if not self.check_contact(right_g_group, o_geoms): 619 | for left_g_group in left_g_geoms: 620 | if not self.check_contact(left_g_group, o_geoms): 621 | return False 622 | return True 623 | 624 | def _post_action(self, action): 625 | """ 626 | Extend post action house keeping to add Success to info. 627 | Args: 628 | action (np.array): Action to execute within the environment 629 | Returns: 630 | 3-tuple: 631 | - (float) reward from the environment 632 | - (bool) whether the current episode is completed or not 633 | - (dict) info dict containing Success state 634 | """ 635 | reward, _, info = super()._post_action(action) 636 | info['Success'] = 1 if reward == 1 else 0 637 | 638 | return reward, self.done, info 639 | -------------------------------------------------------------------------------- /composuite/env/gym_wrapper.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from gym import spaces 3 | 4 | from robosuite.wrappers.gym_wrapper import GymWrapper 5 | 6 | class GymWrapper(GymWrapper): 7 | """Extends the robosuite Gym wrapper which mimics many of the 8 | required functionalities of the Wrapper class found in the gym.core module. 9 | """ 10 | 11 | def __init__(self, env, keys=None): 12 | """Args: 13 | env (MujocoEnv): The environment to wrap. 14 | keys (None, list of str): If provided, each observation will 15 | consist of concatenated keys from the wrapped environment's 16 | observation dictionary. If None, all the keys from the 17 | compositional environment will be added. 18 | """ 19 | # Run super method 20 | super().__init__(env=env) 21 | # Create name for gym 22 | robots = "".join([type(robot.robot_model).__name__ for robot in self.env.robots]) 23 | self.name = robots + "_" + type(self.env).__name__ 24 | 25 | # Get reward range 26 | self.reward_range = (0, self.env.reward_scale) 27 | 28 | if keys is None: 29 | keys = [] 30 | # Add object obs if requested 31 | if self.env.use_object_obs: 32 | keys += ["object-state"] 33 | # Add obstacle obs if requested 34 | if hasattr(self.env, "use_obstacle_obs") and self.env.use_obstacle_obs: 35 | keys += ["obstacle-state"] 36 | # Add goal obs if requested 37 | if hasattr(self.env, "use_goal_obs") and self.env.use_goal_obs: 38 | keys += ["goal-state"] 39 | # Add image obs if requested 40 | if self.env.use_camera_obs: 41 | keys += [f"{cam_name}_image" for cam_name in self.env.camera_names] 42 | # Add task obs if requested 43 | if self.env.use_task_id_obs: 44 | # keys += ["task-state"] 45 | keys += ["object_id", "robot_id", "obstacle_id", "subtask_id"] 46 | # Iterate over all robots to add to state 47 | for idx in range(len(self.env.robots)): 48 | keys += ["robot{}_proprio-state".format(idx)] 49 | self.keys = keys 50 | 51 | # Gym specific attributes 52 | self.env.spec = None 53 | self.metadata = None 54 | 55 | # set up observation and action spaces 56 | obs = self.env.reset() 57 | self.modality_dims = {key: obs[key].shape for key in self.keys} 58 | self.observation_positions = {} 59 | idx = 0 60 | for key in self.keys: 61 | self.observation_positions[key] = np.arange(idx, obs[key].shape[0] + idx) 62 | idx += obs[key].shape[0] 63 | 64 | flat_ob = self._flatten_obs(obs) 65 | self.obs_dim = flat_ob.size 66 | high = np.inf * np.ones(self.obs_dim) 67 | low = -high 68 | self.observation_space = spaces.Box(low=low, high=high) 69 | low, high = self.env.action_spec 70 | self.action_space = spaces.Box(low=low, high=high) 71 | -------------------------------------------------------------------------------- /composuite/env/main.py: -------------------------------------------------------------------------------- 1 | from cgitb import small 2 | import warnings 3 | import itertools 4 | import random 5 | 6 | from robosuite.controllers import load_controller_config 7 | from composuite.env.gym_wrapper import GymWrapper 8 | import robosuite as suite 9 | 10 | AVAILABLE_ROBOTS = ["IIWA", "Jaco", "Kinova3", "Panda"] 11 | AVAILABLE_OBSTACLES = [None, "None", "GoalWall", 12 | "ObjectDoor", "ObjectWall"] 13 | AVAILABLE_OBJECTS = ["Box", "Dumbbell", "Plate", "Hollowbox"] 14 | AVAILABLE_TASKS = ["PickPlace", "Push", "Shelf", "Trashcan"] 15 | 16 | 17 | def make(robot="IIWA", obj="Box", obstacle=None, 18 | task="PickPlace", controller="joint", env_horizon=500, 19 | has_renderer=False, has_offscreen_renderer=False, 20 | reward_shaping=True, ignore_done=True, use_camera_obs=False, 21 | **kwargs) -> GymWrapper: 22 | """Create a compositional environment in form of a gym environment. 23 | 24 | Args: 25 | robot (str, optional): Robot to use in environment. 26 | Defaults to "IIWA". 27 | obj (str, optional): Object to use in environment. 28 | Defaults to "milk". 29 | obstacle (_type_, optional): Obstacle to use in environment. 30 | Defaults to None. 31 | task (str, optional): Objective to use in environment. 32 | Defaults to "PickPlace". 33 | controller (str, optional): Robot controller. 34 | Options are osc, joint 35 | and osc_pose. Defaults to "osc". 36 | env_horizon (int, optional): Number of steps after which env resets. 37 | Defaults to 500. 38 | has_renderer (bool, optional): True, if environment should have a renderer 39 | to support visualization. 40 | Defaults to False. 41 | has_offscreen_renderer (bool, optional): True, if environment 42 | should have an offscreen renderer 43 | to support visual observations. 44 | Defaults to False. 45 | reward_shaping (bool, optional): True, if shaped rewards instead of sparse 46 | rewards should be used. Defaults to True. 47 | ignore_done (bool, optional): True, if environment should not output done 48 | after reaching horizon. Defaults to True. 49 | use_camera_obs (bool, optional): True, if environment should return visual 50 | observations. Defaults to False. 51 | 52 | Raises: 53 | NotImplementedError: Raises if environment name is misspecified. 54 | NotImplementedError: Raises if controller is misspecified. 55 | 56 | Returns: 57 | GymWrapper: Gym-like environment 58 | """ 59 | 60 | assert robot in AVAILABLE_ROBOTS 61 | assert obstacle in AVAILABLE_OBSTACLES 62 | assert obj in AVAILABLE_OBJECTS 63 | assert task in AVAILABLE_TASKS 64 | 65 | if obstacle == "None": 66 | obstacle = None 67 | 68 | # defined options to create robosuite environment 69 | options = {} 70 | 71 | if task == "PickPlace": 72 | options["env_name"] = "PickPlaceSubtask" 73 | elif task == "Push": 74 | options["env_name"] = "PushSubtask" 75 | elif task == "Shelf": 76 | options["env_name"] = "ShelfSubtask" 77 | elif task == "Trashcan": 78 | options["env_name"] = "TrashcanSubtask" 79 | else: 80 | raise NotImplementedError 81 | 82 | options["robots"] = robot 83 | options["obstacle"] = obstacle 84 | options["object_type"] = obj 85 | 86 | if controller == "osc": 87 | controller_name = "OSC_POSITION" 88 | elif controller == "joint": 89 | controller_name = "JOINT_POSITION" 90 | elif controller == 'osc_pose': 91 | controller_name = 'OSC_POSE' 92 | else: 93 | print("Controller unknown") 94 | raise NotImplementedError 95 | 96 | options["controller_configs"] = load_controller_config( 97 | default_controller=controller_name) 98 | 99 | env = suite.make( 100 | **options, 101 | has_renderer=has_renderer, 102 | has_offscreen_renderer=has_offscreen_renderer, 103 | reward_shaping=reward_shaping, 104 | ignore_done=ignore_done, 105 | use_camera_obs=use_camera_obs, 106 | horizon=env_horizon, 107 | **kwargs 108 | ) 109 | 110 | env.reset() 111 | return GymWrapper(env) 112 | 113 | 114 | def sample_tasks(experiment_type='default', num_train=1, 115 | smallscale_elem=None, holdout_elem=None, 116 | shuffling_seed=None, no_shuffle=False): 117 | """Sample a set of training and test configurations as tasks. 118 | 119 | Args: 120 | experiment_type (str, optional): The type of experiment that 121 | is used for sampling. The possible experiments are the experiments 122 | referred to in the initial publication. 123 | Options are: default, smallscale, holdout. Defaults to 'default'. 124 | num_train (int, optional): Number of training tasks to sample. 125 | All other tasks will be returned as test tasks. Defaults to 1. 126 | smallscale_elem (_type_, optional): Required when experiment_type 127 | smallscale. The axis element that is to be fixed. No other element 128 | from that axis will be sampled. Defaults to None. 129 | holdout_elem (_type_, optional): Required when experiment_type is 130 | holdout. Only a single configuration containing this element will 131 | be returned in the train set. The test set contains all other tasks 132 | using this element other than the single one in the train set. 133 | Defaults to None. 134 | shuffling_seed (_type_, optional): Random seed for shuffling configurations. 135 | This will reset to the state before shuffling so that your random 136 | seed selection outside this method is not affected. 137 | Defaults to None. 138 | 139 | Raises: 140 | NotImplementedError: Raises if experiment type is unkown. 141 | NotImplementedError: Raises if selected 142 | axis element in non-default experiments is unknown. 143 | 144 | Returns: 145 | List[Tuple]: List containing the training task configurations. 146 | List[Tuple]: List containing the test task configurations. 147 | """ 148 | if experiment_type == 'smallscale': 149 | assert num_train >= 0 and num_train <= 64, \ 150 | "Number of tasks must be at least 1 and at \ 151 | most 64 in smallscale setting." 152 | elif experiment_type == 'holdout': 153 | assert num_train >= 1 and num_train <= 193, \ 154 | "Number of tasks must be at least 1 and at \ 155 | most 193 in holdout setting." 156 | elif experiment_type == 'default': 157 | assert num_train >= 0 and num_train <= 256, \ 158 | "Number of tasks must be at least 1 and at \ 159 | most 256 in default setting" 160 | else: 161 | raise NotImplementedError("Specified experiment type does not exist. \ 162 | Options are: default, smallscale, holdout.") 163 | 164 | if experiment_type == 'default': 165 | if smallscale_elem is not None or holdout_elem is not None: 166 | warnings.warn("You specified a smallscale/holdout element but are \ 167 | using the default experiment setting. Continuing and ignoring \ 168 | smallscale/holdout element.") 169 | elif experiment_type == 'smallscale': 170 | assert smallscale_elem is not None, \ 171 | "Selected experiment type smallscale \ 172 | but did not specifiy fixed element." 173 | elif experiment_type == 'holdout': 174 | assert holdout_elem is not None, \ 175 | "Selected experiment type holdout \ 176 | but did not specifiy single task element." 177 | 178 | if experiment_type == 'default' or experiment_type == 'holdout': 179 | _robots = AVAILABLE_ROBOTS 180 | _objects = AVAILABLE_OBJECTS 181 | _obstacles = [e for e in AVAILABLE_OBSTACLES if e is not None] 182 | _objectives = AVAILABLE_TASKS 183 | elif experiment_type == 'smallscale': 184 | _robots = AVAILABLE_ROBOTS 185 | _objects = AVAILABLE_OBJECTS 186 | _obstacles = [e for e in AVAILABLE_OBSTACLES if e is not None] 187 | _objectives = AVAILABLE_TASKS 188 | 189 | if smallscale_elem in AVAILABLE_ROBOTS: 190 | _robots = [smallscale_elem] 191 | elif smallscale_elem in AVAILABLE_OBJECTS: 192 | _objects = [smallscale_elem] 193 | elif smallscale_elem in AVAILABLE_OBSTACLES: 194 | _obstacles = [smallscale_elem] 195 | elif smallscale_elem in AVAILABLE_TASKS: 196 | _objectives = [smallscale_elem] 197 | else: 198 | raise NotImplementedError( 199 | "Specified smallscale element does not exist. Options are: \n\ 200 | IIWA, Jaco, Kinova3, Panda, None, GoalWall, \ 201 | ObjectDoor, ObjectWall, Box, Dumbbell, Plate, Hollowbox \ 202 | PickPlace, Push, Shelf, Trashcan") 203 | 204 | all_configurations = [] 205 | for conf in itertools.product(_robots, _objects, _obstacles, _objectives): 206 | all_configurations.append(conf) 207 | 208 | state = random.getstate() 209 | if shuffling_seed is not None: 210 | random.seed(shuffling_seed) 211 | if not no_shuffle: 212 | random.shuffle(all_configurations) 213 | random.setstate(state) 214 | 215 | if experiment_type == 'default' or experiment_type == 'smallscale': 216 | return all_configurations[:num_train], all_configurations[num_train:] 217 | elif experiment_type == 'holdout': 218 | elem_tasks = [ 219 | conf for conf in all_configurations if holdout_elem in conf] 220 | non_elem_tasks = [ 221 | conf for conf in all_configurations if holdout_elem not in conf] 222 | 223 | train_elem_task = elem_tasks[:1] 224 | holdout_tasks = elem_tasks[1:] 225 | 226 | return non_elem_tasks[:num_train - 1] + train_elem_task, holdout_tasks 227 | -------------------------------------------------------------------------------- /composuite/objects/objects.py: -------------------------------------------------------------------------------- 1 | from robosuite.models.objects import MujocoXMLObject 2 | from composuite.utils.mjk_utils import xml_path_completion 3 | 4 | import numpy as np 5 | 6 | 7 | class ObjectWallObject(MujocoXMLObject): 8 | """ 9 | Wall obstacle (used in CompositionalEnv) 10 | """ 11 | 12 | def __init__(self, name): 13 | super().__init__(xml_path_completion("objects/object_wall.xml"), 14 | name=name, obj_type="all", joints=None, duplicate_collision_geoms=True) 15 | 16 | self.wall_size = np.array(list(map(float, self.worldbody.findall( 17 | "./body/body/geom")[0].items()[3][1].split()))) 18 | 19 | 20 | class ObjectDoorFrameObject(MujocoXMLObject): 21 | """ 22 | DoorFrame obstacle (used in CompositionalEnv) 23 | """ 24 | 25 | def __init__(self, name): 26 | super().__init__(xml_path_completion("objects/object_door_frame.xml"), 27 | name=name, obj_type="all", joints=None, duplicate_collision_geoms=True) 28 | 29 | self.door_l_size = np.array(list(map(float, self.worldbody.findall( 30 | "./body/body/body/geom")[0].items()[3][1].split()))) 31 | 32 | 33 | class GoalWallObject(MujocoXMLObject): 34 | """ 35 | Wall obstacle (used in CompositionalEnv) 36 | """ 37 | 38 | def __init__(self, name): 39 | super().__init__(xml_path_completion("objects/goal_wall.xml"), 40 | name=name, obj_type="all", joints=None, duplicate_collision_geoms=True) 41 | 42 | 43 | class DumbbellObject(MujocoXMLObject): 44 | """ 45 | Cup object (used in CompositionalEnv) 46 | """ 47 | 48 | def __init__(self, name): 49 | super().__init__(xml_path_completion("objects/dumbbell.xml"), 50 | name=name, obj_type="all", joints=[dict(type="free", damping="0.0005")], duplicate_collision_geoms=True) 51 | 52 | self.upper_radius = np.array(list(map(float, self.worldbody.findall( 53 | "./body/body/geom")[2].items()[1][1].split())))[0] 54 | self.lower_radius = np.array(list(map(float, self.worldbody.findall( 55 | "./body/body/geom")[0].items()[1][1].split())))[0] 56 | 57 | 58 | class DumbbellVisualObject(MujocoXMLObject): 59 | """ 60 | Cup Visual object (used in CompositionalEnv) 61 | """ 62 | 63 | def __init__(self, name): 64 | super().__init__(xml_path_completion("objects/dumbbell-visual.xml"), 65 | name=name, joints=None, obj_type="visual", duplicate_collision_geoms=True) 66 | 67 | 68 | class PlateObject(MujocoXMLObject): 69 | """ 70 | Plate object (used in CompositionalEnv) 71 | """ 72 | 73 | def __init__(self, name): 74 | super().__init__(xml_path_completion("objects/plate.xml"), 75 | name=name, obj_type="all", joints=[dict(type="free", damping="0.0005")], duplicate_collision_geoms=True) 76 | 77 | self.radius = np.array(list(map(float, self.worldbody.findall( 78 | "./body/body/geom")[2].items()[1][1].split())))[0] 79 | 80 | 81 | class PlateVisualObject(MujocoXMLObject): 82 | """ 83 | Plate Visual object (used in CompositionalEnv) 84 | """ 85 | 86 | def __init__(self, name): 87 | super().__init__(xml_path_completion("objects/plate-visual.xml"), 88 | name=name, joints=None, obj_type="visual", duplicate_collision_geoms=True) 89 | 90 | 91 | class HollowBoxObject(MujocoXMLObject): 92 | """ 93 | Hollow box object (used in CompositionalEnv) 94 | """ 95 | 96 | def __init__(self, name): 97 | super().__init__(xml_path_completion("objects/hollowbox.xml"), 98 | name=name, obj_type="all", joints=[dict(type="free", damping="0.0005")], duplicate_collision_geoms=True) 99 | 100 | self.length = np.array(list(map(float, self.worldbody.findall( 101 | "./body/body/geom")[0].items()[1][1].split())))[0] 102 | 103 | 104 | class HollowBoxVisualObject(MujocoXMLObject): 105 | """ 106 | Hollow box Visual object (used in CompositionalEnv) 107 | """ 108 | 109 | def __init__(self, name): 110 | super().__init__(xml_path_completion("objects/hollowbox-visual.xml"), 111 | name=name, joints=None, obj_type="visual", duplicate_collision_geoms=True) 112 | 113 | 114 | class CustomBoxObject(MujocoXMLObject): 115 | """ 116 | Hollow box object (used in CompositionalEnv) 117 | """ 118 | 119 | def __init__(self, name): 120 | super().__init__(xml_path_completion("objects/custombox.xml"), 121 | name=name, obj_type="all", joints=[dict(type="free", damping="0.0005")], duplicate_collision_geoms=True) 122 | 123 | self.length = np.array(list(map(float, self.worldbody.findall( 124 | "./body/body/geom")[0].items()[1][1].split())))[0] 125 | 126 | 127 | class CustomBoxVisualObject(MujocoXMLObject): 128 | """ 129 | Hollow box object (used in CompositionalEnv) 130 | """ 131 | 132 | def __init__(self, name): 133 | super().__init__(xml_path_completion("objects/custombox-visual.xml"), 134 | name=name, joints=None, obj_type="visual", duplicate_collision_geoms=True) 135 | 136 | self.length = np.array(list(map(float, self.worldbody.findall( 137 | "./body/body/geom")[0].items()[1][1].split())))[0] 138 | -------------------------------------------------------------------------------- /composuite/tasks/pick_place_subtask.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from composuite.arenas.pick_place_arena import PickPlaceArena 4 | from composuite.env.compositional_env import CompositionalEnv 5 | 6 | import robosuite.utils.transform_utils as T 7 | from robosuite.utils.placement_samplers import UniformRandomSampler 8 | 9 | 10 | class PickPlaceSubtask(CompositionalEnv): 11 | """This class corresponds to the pick place task for a single robot arm. 12 | """ 13 | 14 | def __init__( 15 | self, 16 | robots, 17 | object_type, 18 | obstacle, 19 | env_configuration="default", 20 | controller_configs=None, 21 | mount_types="default", 22 | gripper_types="RethinkGripper", 23 | initialization_noise=None, 24 | use_camera_obs=True, 25 | use_object_obs=True, 26 | use_task_id_obs=False, 27 | has_renderer=False, 28 | has_offscreen_renderer=True, 29 | render_camera="frontview", 30 | render_collision_mesh=False, 31 | render_visual_mesh=True, 32 | render_gpu_device_id=-1, 33 | control_freq=20, 34 | horizon=1000, 35 | ignore_done=True, 36 | hard_reset=True, 37 | camera_names="agentview", 38 | camera_heights=256, 39 | camera_widths=256, 40 | camera_depths=False, 41 | bin1_pos=(0.1, -0.26, 0.8), 42 | bin2_pos=(0.1, 0.13, 0.8), 43 | reward_scale=1.0, 44 | reward_shaping=False, 45 | ): 46 | 47 | self.subtask_id = 0 48 | super().__init__( 49 | robots, 50 | object_type, 51 | obstacle, 52 | bin1_pos, 53 | bin2_pos, 54 | env_configuration=env_configuration, 55 | controller_configs=controller_configs, 56 | mount_types=mount_types, 57 | gripper_types=gripper_types, 58 | initialization_noise=initialization_noise, 59 | use_camera_obs=use_camera_obs, 60 | use_object_obs=use_object_obs, 61 | use_task_id_obs=use_task_id_obs, 62 | has_renderer=has_renderer, 63 | has_offscreen_renderer=has_offscreen_renderer, 64 | render_camera=render_camera, 65 | render_collision_mesh=render_collision_mesh, 66 | render_visual_mesh=render_visual_mesh, 67 | render_gpu_device_id=render_gpu_device_id, 68 | control_freq=control_freq, 69 | horizon=horizon, 70 | ignore_done=ignore_done, 71 | hard_reset=hard_reset, 72 | camera_names=camera_names, 73 | camera_heights=camera_heights, 74 | camera_widths=camera_widths, 75 | camera_depths=camera_depths, 76 | reward_scale=reward_scale, 77 | reward_shaping=reward_shaping, 78 | ) 79 | 80 | def staged_rewards(self, action): 81 | """ 82 | Returns staged rewards based on current physical states. 83 | Stages consist of reaching, grasping, lifting, and hovering. 84 | 85 | Returns: 86 | - (float) reaching reward 87 | - (float) grasping reward 88 | - (float) lifting reward 89 | - (float) hovering reward 90 | """ 91 | 92 | reach_mult = 0.2 93 | grasp_mult = 0.3 94 | lift_mult = 0.5 95 | hover_mult = 0.7 96 | drop_mult = 0.9 97 | 98 | # reaching reward governed by distance to closest object 99 | r_reach = 0. 100 | if not self.object_in_bin: 101 | # get reaching reward via minimum distance to a target object 102 | dist = self._gripper_to_target( 103 | gripper=self.robots[0].gripper, 104 | target=self.object.root_body, 105 | target_type="body", 106 | return_distance=True, 107 | ) 108 | 109 | r_reach = (1 - np.tanh(10.0 * dist)) * reach_mult 110 | 111 | # grasping reward for touching any objects of interest 112 | is_grasping = self._check_grasp( 113 | gripper=self.robots[0].gripper, 114 | object_geoms=[g for g in self.object.contact_geoms]) 115 | r_grasp = int(is_grasping) * grasp_mult 116 | 117 | # lifting reward for picking up an object 118 | r_lift = 0. 119 | if not self.object_in_bin and r_grasp > 0.: 120 | z_target = self.bin2_pos[2] + 0.25 121 | object_z_loc = self.sim.data.body_xpos[self.obj_body_id, 2] 122 | z_dist = np.abs(z_target - object_z_loc) 123 | r_lift = grasp_mult + (1 - np.tanh(5.0 * z_dist)) * ( 124 | lift_mult - grasp_mult 125 | ) 126 | 127 | # segment objects into left of the bins and above the bins 128 | object_xy_loc = self.sim.data.body_xpos[self.obj_body_id, :2] 129 | y_check = ( 130 | np.abs(object_xy_loc[1] - 131 | self.bin2_pos[1]) 132 | < self.bin2_size[1] 133 | ) 134 | x_check = ( 135 | np.abs(object_xy_loc[0] - 136 | self.bin2_pos[0]) 137 | < self.bin2_size[0] 138 | ) 139 | object_above_bin = x_check and y_check 140 | 141 | # hover reward for getting object above bin 142 | r_hover = 0. 143 | r_drop = 0. 144 | if not self.object_in_bin and r_lift > 0.45: 145 | dist = np.linalg.norm( 146 | self.bin2_pos[:2] - object_xy_loc 147 | ) 148 | # objects to the left get r_lift added to hover reward, 149 | # those on the right get max(r_lift) added (to encourage dropping) 150 | if not object_above_bin: 151 | r_hover = r_lift + ( 152 | 1 - np.tanh(2.0 * dist) 153 | ) * (hover_mult - lift_mult) 154 | else: 155 | r_hover = lift_mult + ( 156 | 1 - np.tanh(2.0 * dist) 157 | ) * (hover_mult - lift_mult) 158 | 159 | if r_grasp > 0 and object_above_bin: 160 | z_target = self.bin2_pos[2] + 0.1 161 | object_z_loc = self.sim.data.body_xpos[self.obj_body_id, 2] 162 | z_dist = np.maximum(object_z_loc - z_target, 0.) 163 | r_drop = hover_mult + \ 164 | (1 - np.tanh(5.0 * z_dist)) * (drop_mult - hover_mult) 165 | 166 | return r_reach, r_grasp, r_lift, r_hover, r_drop 167 | 168 | def not_in_bin(self, obj_pos): 169 | """Checks whether object is in the second bin. 170 | 171 | Args: 172 | obj_pos (np.array): Current position of the object. 173 | 174 | Returns: 175 | bool: True if the object is NOT inside the second bin. 176 | """ 177 | bin_x_low = self.bin2_pos[0] - self.bin2_size[0] 178 | bin_y_low = self.bin2_pos[1] - self.bin2_size[1] 179 | 180 | bin_x_high = self.bin2_pos[0] + self.bin2_size[0] 181 | bin_y_high = self.bin2_pos[1] + self.bin2_size[1] 182 | 183 | res = True 184 | if ( 185 | bin_x_low < obj_pos[0] < bin_x_high 186 | and bin_y_low < obj_pos[1] < bin_y_high 187 | and self.bin2_pos[2] < obj_pos[2] < self.bin2_pos[2] + 0.1 188 | ): 189 | res = False 190 | return res 191 | 192 | def _get_placement_initializer(self): 193 | """Helper function for defining placement initializer 194 | and object sampling bounds. Extends super class initilizer. 195 | """ 196 | super()._get_placement_initializer() 197 | 198 | # TODO: why is this not exactly in the middle 199 | self.placement_initializer.append_sampler( 200 | sampler=UniformRandomSampler( 201 | name=f"{self.visual_object.name}ObjectSampler", 202 | mujoco_objects=self.visual_object, 203 | x_range=[0, 0], 204 | y_range=[0, 0], 205 | rotation=0., 206 | rotation_axis='z', 207 | ensure_object_boundary_in_range=False, 208 | ensure_valid_placement=False, 209 | reference_pos=self.bin2_pos, 210 | z_offset=self.bin2_pos[2] - self.bin1_pos[2], 211 | ) 212 | ) 213 | 214 | def _load_model(self): 215 | """ 216 | Loads an xml model, puts it in self.model 217 | """ 218 | 219 | # load model for table top workspace 220 | self.mujoco_arena = PickPlaceArena( 221 | bin1_pos=self.bin1_pos, 222 | ) 223 | 224 | # Load model propagation 225 | super()._load_model() 226 | 227 | # Generate placement initializer 228 | self._initialize_model() 229 | self._get_placement_initializer() 230 | 231 | def _setup_references(self): 232 | """ 233 | Sets up references to important components. A reference is typically an 234 | index or a list of indices that point to the corresponding elements 235 | in a flatten array, which is how MuJoCo stores physical simulation data. 236 | """ 237 | 238 | # keep track of which objects are in their corresponding bins 239 | self.object_in_bin = False 240 | 241 | # target locations in bin for each object type 242 | self.target_bin_placements = np.zeros((1, 3)) 243 | 244 | # TODO: fix this once i understand why its here 245 | # I think we can remove target bin placements 246 | bin_x_low = self.bin2_pos[0] 247 | bin_y_low = self.bin2_pos[1] 248 | bin_x_low += self.bin2_size[0] / 2. 249 | bin_y_low += self.bin2_size[1] / 2. 250 | self.target_bin_placements[0, :] = [ 251 | bin_x_low, bin_y_low, self.bin2_pos[2]] 252 | 253 | super()._setup_references() 254 | 255 | def _reset_internal(self): 256 | """Resets simulation internal configurations. 257 | """ 258 | super()._reset_internal() 259 | 260 | # Set the bins to the desired position 261 | self.sim.model.body_pos[self.sim.model.body_name2id( 262 | "bin1")] = self.bin1_pos 263 | self.sim.model.body_pos[self.sim.model.body_name2id( 264 | "bin2")] = self.bin2_pos 265 | 266 | def _check_success(self): 267 | """ 268 | Check if all objects have been successfully placed in their corresponding bins. 269 | 270 | Returns: 271 | bool: True if object is placed correctly 272 | """ 273 | # remember objects that are in the correct bins 274 | gripper_site_pos = self.sim.data.site_xpos[self.robots[0].eef_site_id] 275 | obj_pos = self.sim.data.body_xpos[self.obj_body_id] 276 | dist = np.linalg.norm(gripper_site_pos - obj_pos) 277 | r_reach = 1 - np.tanh(10.0 * dist) 278 | # self.object_in_bin = not self.not_in_bin(obj_pos) 279 | self.object_in_bin = bool( 280 | (not self.not_in_bin(obj_pos)) and r_reach > 0.35) 281 | 282 | return self.object_in_bin 283 | 284 | -------------------------------------------------------------------------------- /composuite/tasks/push_subtask.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from composuite.arenas.push_arena import PushArena 4 | from composuite.env.compositional_env import CompositionalEnv 5 | 6 | from robosuite.models.tasks import ManipulationTask 7 | from robosuite.utils.placement_samplers import SequentialCompositeSampler, UniformRandomSampler 8 | 9 | 10 | class PushSubtask(CompositionalEnv): 11 | """This class corresponds to the lifting task for a single robot arm. 12 | """ 13 | 14 | def __init__( 15 | self, 16 | robots, 17 | object_type, 18 | obstacle, 19 | env_configuration="default", 20 | controller_configs=None, 21 | mount_types="default", 22 | gripper_types="RethinkGripper", 23 | initialization_noise=None, 24 | use_camera_obs=True, 25 | use_object_obs=True, 26 | use_task_id_obs=False, 27 | has_renderer=False, 28 | has_offscreen_renderer=True, 29 | render_camera="frontview", 30 | render_collision_mesh=False, 31 | render_visual_mesh=True, 32 | render_gpu_device_id=-1, 33 | control_freq=20, 34 | horizon=1000, 35 | ignore_done=True, 36 | hard_reset=True, 37 | camera_names="agentview", 38 | camera_heights=256, 39 | camera_widths=256, 40 | camera_depths=False, 41 | bin1_pos=(0.1, -0.26, 0.8), 42 | bin2_pos=(0.1, 0.13, 0.8), 43 | reward_scale=1.0, 44 | reward_shaping=False, 45 | ): 46 | 47 | self.subtask_id = 1 48 | super().__init__( 49 | robots, 50 | object_type, 51 | obstacle, 52 | bin1_pos, 53 | bin2_pos, 54 | env_configuration=env_configuration, 55 | controller_configs=controller_configs, 56 | mount_types=mount_types, 57 | gripper_types=gripper_types, 58 | initialization_noise=initialization_noise, 59 | use_camera_obs=use_camera_obs, 60 | use_object_obs=use_object_obs, 61 | use_task_id_obs=use_task_id_obs, 62 | has_renderer=has_renderer, 63 | has_offscreen_renderer=has_offscreen_renderer, 64 | render_camera=render_camera, 65 | render_collision_mesh=render_collision_mesh, 66 | render_visual_mesh=render_visual_mesh, 67 | render_gpu_device_id=render_gpu_device_id, 68 | control_freq=control_freq, 69 | horizon=horizon, 70 | ignore_done=ignore_done, 71 | hard_reset=hard_reset, 72 | camera_names=camera_names, 73 | camera_heights=camera_heights, 74 | camera_widths=camera_widths, 75 | camera_depths=camera_depths, 76 | reward_scale=reward_scale, 77 | reward_shaping=reward_shaping, 78 | ) 79 | 80 | def staged_rewards(self, action): 81 | """ 82 | Returns staged rewards based on current physical states. 83 | Stages consist of reaching, grasping and approaching. 84 | 85 | Returns: 86 | - (float) reaching reward 87 | - (float) grasping reward 88 | - (float) approach reward 89 | """ 90 | reach_mult = 0.2 91 | grasp_mult = 0.3 92 | approach_mult = 0.7 93 | 94 | # reaching reward governed by distance to object 95 | obj_pos = self.sim.data.body_xpos[self.obj_body_id] 96 | gripper_site_pos = self.sim.data.site_xpos[self.robots[0].eef_site_id] 97 | dist = np.linalg.norm(gripper_site_pos - obj_pos) 98 | r_reach = (1 - np.tanh(10.0 * dist)) * reach_mult 99 | 100 | # touch reward for grasping the object 101 | r_grasp = int(self._check_grasp( 102 | gripper=self.robots[0].gripper, 103 | object_geoms=[g for g in self.object.contact_geoms]) 104 | ) * grasp_mult 105 | 106 | r_approach = 0. 107 | if r_grasp > 0: 108 | # approach reward for approaching goal with object 109 | goal_pos = self.sim.model.body_pos[self.goal_body_id] 110 | dist = np.linalg.norm(goal_pos - obj_pos) 111 | r_approach = grasp_mult + \ 112 | (1 - np.tanh(5. * dist)) * (approach_mult - grasp_mult) 113 | 114 | return r_reach, r_grasp, r_approach 115 | 116 | def _load_model(self): 117 | """Loads an xml model, puts it in self.model 118 | """ 119 | # Load model propagation 120 | 121 | # load model for table top workspace 122 | self.mujoco_arena = PushArena( 123 | bin1_pos=self.bin1_pos, 124 | bin2_pos=self.bin2_pos, 125 | ) 126 | 127 | super()._load_model() 128 | self._initialize_model() 129 | self._get_placement_initializer() 130 | 131 | def _get_placement_initializer(self): 132 | """Helper function for defining placement initializer and object sampling bounds. 133 | """ 134 | super()._get_placement_initializer() 135 | 136 | # placement is relative to object bin, so compute difference and send to placement initializer 137 | self.placement_initializer.append_sampler( 138 | sampler=UniformRandomSampler( 139 | name=f"{self.visual_object.name}ObjectSampler", 140 | mujoco_objects=self.visual_object, 141 | x_range=[0, 0], 142 | y_range=[0, 0], 143 | rotation=0., 144 | rotation_axis='z', 145 | ensure_object_boundary_in_range=False, 146 | ensure_valid_placement=False, 147 | reference_pos=self.bin2_pos, 148 | z_offset=0., 149 | ) 150 | ) 151 | 152 | def _setup_references(self): 153 | """Sets up references to important components. A reference is typically an 154 | index or a list of indices that point to the corresponding elements 155 | in a flatten array, which is how MuJoCo stores physical simulation data. 156 | """ 157 | 158 | super()._setup_references() 159 | 160 | def _reset_internal(self): 161 | """Resets simulation internal configurations. 162 | """ 163 | super()._reset_internal() 164 | 165 | def _check_success(self): 166 | """Check if cube has been lifted. 167 | 168 | Returns: 169 | bool: True if cube has been lifted 170 | """ 171 | obj_xy_pos = self.sim.data.body_xpos[self.obj_body_id][:2] 172 | goal_xy_pos = self.sim.model.body_pos[self.goal_body_id][:2] 173 | 174 | obj_height = self.sim.data.body_xpos[self.obj_body_id][2] 175 | table_height = self.bin2_pos[2] 176 | 177 | return np.linalg.norm(obj_xy_pos - goal_xy_pos) <= 0.03 and obj_height <= table_height + 0.1 178 | 179 | def _post_action(self, action): 180 | """ 181 | Do any housekeeping after taking an action. 182 | 183 | Args: 184 | action (np.array): Action to execute within the environment 185 | 186 | Returns: 187 | - (float) reward from the environment 188 | - (bool) whether the current episode is completed or not 189 | - (dict) empty dict to be filled with information by subclassed method 190 | 191 | """ 192 | reward, _, info = super()._post_action(action) 193 | 194 | object_height = self.sim.data.body_xpos[self.obj_body_id][2] 195 | table_height = self.bin2_pos[2] 196 | 197 | self.done = object_height - \ 198 | table_height > 0.1 or ( 199 | (self.timestep >= self.horizon) and not self.ignore_done) 200 | 201 | return reward, self.done, info 202 | -------------------------------------------------------------------------------- /composuite/tasks/shelf_subtask.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from composuite.arenas.shelf_arena import ShelfArena 4 | from composuite.env.compositional_env import CompositionalEnv 5 | 6 | import robosuite.utils.transform_utils as T 7 | from robosuite.utils.placement_samplers import UniformRandomSampler 8 | 9 | 10 | class ShelfSubtask(CompositionalEnv): 11 | """This class corresponds to the shelf task for a single robot arm. 12 | """ 13 | 14 | def __init__( 15 | self, 16 | robots, 17 | object_type, 18 | obstacle, 19 | env_configuration="default", 20 | controller_configs=None, 21 | mount_types="default", 22 | gripper_types="RethinkGripper", 23 | initialization_noise=None, 24 | use_camera_obs=True, 25 | use_object_obs=True, 26 | use_task_id_obs=False, 27 | has_renderer=False, 28 | has_offscreen_renderer=True, 29 | render_camera="frontview", 30 | render_collision_mesh=False, 31 | render_visual_mesh=True, 32 | render_gpu_device_id=-1, 33 | control_freq=20, 34 | horizon=1000, 35 | ignore_done=True, 36 | hard_reset=True, 37 | camera_names="agentview", 38 | camera_heights=256, 39 | camera_widths=256, 40 | camera_depths=False, 41 | bin1_pos=(0.1, -0.26, 0.8), 42 | bin2_pos=(0.1, 0.18, 0.8), 43 | reward_scale=1.0, 44 | reward_shaping=False, 45 | ): 46 | 47 | self.subtask_id = 2 48 | super().__init__( 49 | robots, 50 | object_type, 51 | obstacle, 52 | bin1_pos, 53 | bin2_pos, 54 | env_configuration=env_configuration, 55 | controller_configs=controller_configs, 56 | mount_types=mount_types, 57 | gripper_types=gripper_types, 58 | initialization_noise=initialization_noise, 59 | use_camera_obs=use_camera_obs, 60 | use_object_obs=use_object_obs, 61 | use_task_id_obs=use_task_id_obs, 62 | has_renderer=has_renderer, 63 | has_offscreen_renderer=has_offscreen_renderer, 64 | render_camera=render_camera, 65 | render_collision_mesh=render_collision_mesh, 66 | render_visual_mesh=render_visual_mesh, 67 | render_gpu_device_id=render_gpu_device_id, 68 | control_freq=control_freq, 69 | horizon=horizon, 70 | ignore_done=ignore_done, 71 | hard_reset=hard_reset, 72 | camera_names=camera_names, 73 | camera_heights=camera_heights, 74 | camera_widths=camera_widths, 75 | camera_depths=camera_depths, 76 | reward_scale=reward_scale, 77 | reward_shaping=reward_shaping, 78 | ) 79 | 80 | def staged_rewards(self, action): 81 | """ 82 | Returns staged rewards based on current physical states. 83 | Stages consist of reaching, grasping, lifting, aligning and approaching. 84 | 85 | Returns: 86 | - (float) reaching reward 87 | - (float) grasping reward 88 | - (float) lifting reward 89 | - (float) aligning reward 90 | - (float) approaching reward 91 | """ 92 | 93 | reach_mult = 0.2 94 | grasp_mult = 0.3 95 | lift_mult = 0.5 96 | align_mult = 0.8 97 | approach_mult = 0.9 98 | 99 | # reaching reward governed by distance to closest object 100 | r_reach = 0. 101 | if not self.object_in_shelf: 102 | # get reaching reward via minimum distance to a target object 103 | dist = self._gripper_to_target( 104 | gripper=self.robots[0].gripper, 105 | target=self.object.root_body, 106 | target_type="body", 107 | return_distance=True, 108 | ) 109 | r_reach = (1 - np.tanh(10.0 * dist)) * reach_mult 110 | 111 | # grasping reward for touching any objects of interest 112 | r_grasp = int(self._check_grasp( 113 | gripper=self.robots[0].gripper, 114 | object_geoms=[g for g in self.object.contact_geoms]) 115 | ) * grasp_mult 116 | 117 | object_x_loc = self.sim.data.body_xpos[self.obj_body_id, 0] 118 | shelf_x_low = self.shelf_pos[0] - self.shelf_x_size 119 | shelf_x_high = self.shelf_pos[0] + self.shelf_x_size 120 | x_check = shelf_x_low < object_x_loc < shelf_x_high 121 | 122 | object_z_loc = self.sim.data.body_xpos[self.obj_body_id, 2] 123 | shelf_z_low = self.shelf_z_offset + self.shelf_pos[2] 124 | shelf_z_high = shelf_z_low + self.shelf_z_height 125 | z_check = shelf_z_low + self.mujoco_arena.shelf_thickness < object_z_loc < shelf_z_high 126 | 127 | object_infront_shelf = z_check and x_check 128 | 129 | # lifting reward for picking up an object 130 | r_lift = 0. 131 | if r_grasp > 0.: 132 | z_target = self.bin2_pos[2] + 0.25 133 | object_z_loc = self.sim.data.body_xpos[self.obj_body_id, 2] 134 | dist = np.abs(z_target - object_z_loc) 135 | r_lift = r_grasp + (1 - np.tanh(5.0 * dist)) * ( 136 | lift_mult - grasp_mult 137 | ) 138 | 139 | r_align = 0 140 | # align reward 141 | if object_infront_shelf: 142 | gripper_quat = T.convert_quat(self.sim.data.get_body_xquat( 143 | self.robots[0].robot_model.eef_name), to="xyzw") 144 | gripper_mat = T.quat2mat(gripper_quat) 145 | 146 | gripper_in_world = gripper_mat.dot(np.array([0, 0, 1])) 147 | target_orientation = np.array([0, 1, 0]) 148 | orientation_dist = self._dot_product_angle( 149 | gripper_in_world, target_orientation) 150 | 151 | r_align = lift_mult + \ 152 | (1 - np.tanh(orientation_dist)) * \ 153 | (align_mult - lift_mult) 154 | 155 | r_approach = 0 156 | if object_infront_shelf and r_align > 0.6: 157 | # approach 158 | y_target = self.shelf_pos[1] - self.shelf_y_size 159 | y_dist = np.maximum( 160 | y_target - self.sim.data.body_xpos[self.obj_body_id, 1], 0.) 161 | r_approach = align_mult + \ 162 | (1 - np.tanh(5.0 * y_dist)) * \ 163 | (approach_mult - align_mult) 164 | 165 | return r_reach, r_grasp, r_lift, r_align, r_approach 166 | 167 | def not_in_shelf(self, obj_pos): 168 | """Checks whether object is inside the shelf. 169 | 170 | Args: 171 | obj_pos (np.array): Current position of the object. 172 | 173 | Returns: 174 | bool: True if the object is NOT inside the shelf. 175 | """ 176 | shelf_x_low = self.shelf_pos[0] - self.shelf_x_size 177 | shelf_x_high = self.shelf_pos[0] + self.shelf_x_size 178 | shelf_y_low = self.shelf_pos[1] - (self.shelf_y_size * 2) * 0.75 179 | shelf_y_high = self.shelf_pos[1] 180 | 181 | res = True 182 | if ( 183 | shelf_x_low < obj_pos[0] < shelf_x_high 184 | and shelf_y_low < obj_pos[1] < shelf_y_high 185 | and self.shelf_z_offset + self.shelf_pos[2] < obj_pos[2] < self.shelf_z_offset + self.shelf_pos[2] + 0.15 186 | ): 187 | res = False 188 | return res 189 | 190 | def _get_placement_initializer(self): 191 | """Helper function for defining placement initializer and object sampling bounds. 192 | """ 193 | super()._get_placement_initializer() 194 | 195 | self.placement_initializer.append_sampler( 196 | sampler=UniformRandomSampler( 197 | name=f"{self.visual_object.name}ObjectSampler", 198 | mujoco_objects=self.visual_object, 199 | x_range=[0, 0], 200 | y_range=[-self.shelf_y_size, -self.shelf_y_size], 201 | rotation=0., 202 | rotation_axis='z', 203 | ensure_object_boundary_in_range=False, 204 | ensure_valid_placement=False, 205 | z_offset=self.shelf_z_offset, 206 | reference_pos=self.shelf_pos, 207 | ) 208 | ) 209 | 210 | def _load_model(self): 211 | """Loads an xml model, puts it in self.model 212 | """ 213 | # load model for table top workspace 214 | self.mujoco_arena = ShelfArena( 215 | bin1_pos=self.bin1_pos, 216 | bin2_pos=self.bin2_pos, 217 | ) 218 | 219 | # Load model propagation 220 | super()._load_model() 221 | 222 | # Generate placement initializer 223 | self._initialize_model() 224 | self._get_placement_initializer() 225 | 226 | def _initialize_model(self): 227 | """Load all the required arena model information from the shelf task 228 | and store it in class variables. 229 | """ 230 | super()._initialize_model() 231 | self.shelf_pos = self.mujoco_arena.shelf_pos 232 | self.shelf_x_size = self.mujoco_arena.shelf_x_size 233 | self.shelf_y_size = self.mujoco_arena.shelf_y_size 234 | self.shelf_z_offset = self.mujoco_arena.shelf_z_offset 235 | self.shelf_z_height = self.mujoco_arena.shelf_z_height 236 | 237 | def _setup_references(self): 238 | """ 239 | Sets up references to important components. A reference is typically an 240 | index or a list of indices that point to the corresponding elements 241 | in a flatten array, which is how MuJoCo stores physical simulation data. 242 | """ 243 | # keep track of if object is in the shelf 244 | self.object_in_shelf = False 245 | super()._setup_references() 246 | 247 | def _check_success(self): 248 | """Check if object has been successfully placed in the shelf. 249 | 250 | Returns: 251 | bool: True if object is placed correctly 252 | """ 253 | # remember if object is in the shelf 254 | obj_pos = self.sim.data.body_xpos[self.obj_body_id] 255 | self.object_in_shelf = not self.not_in_shelf(obj_pos) 256 | 257 | # returns True if the object is in the shelf 258 | return self.object_in_shelf 259 | 260 | def _dot_product_angle(self, v1, v2): 261 | """Computes the dot product angle between two vectors. 262 | """ 263 | if np.linalg.norm(v1) == 0 or np.linalg.norm(v2) == 0: 264 | print("Zero magnitude vector!") 265 | else: 266 | vector_dot_product = np.dot(v1, v2) 267 | arccos = np.arccos(vector_dot_product / 268 | (np.linalg.norm(v1) * np.linalg.norm(v2))) 269 | return arccos 270 | return 0 271 | -------------------------------------------------------------------------------- /composuite/tasks/trashcan_subtask.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from composuite.arenas.trashcan_arena import TrashcanArena 4 | from composuite.env.compositional_env import CompositionalEnv 5 | 6 | from robosuite.utils.placement_samplers import UniformRandomSampler 7 | 8 | 9 | class TrashcanSubtask(CompositionalEnv): 10 | """This class corresponds to the trashcan task for a single robot arm. 11 | """ 12 | 13 | def __init__( 14 | self, 15 | robots, 16 | object_type, 17 | obstacle, 18 | env_configuration="default", 19 | controller_configs=None, 20 | mount_types="default", 21 | gripper_types="RethinkGripper", 22 | initialization_noise=None, 23 | use_camera_obs=True, 24 | use_object_obs=True, 25 | use_task_id_obs=False, 26 | has_renderer=False, 27 | has_offscreen_renderer=True, 28 | render_camera="frontview", 29 | render_collision_mesh=False, 30 | render_visual_mesh=True, 31 | render_gpu_device_id=-1, 32 | control_freq=20, 33 | horizon=1000, 34 | ignore_done=True, 35 | hard_reset=True, 36 | camera_names="agentview", 37 | camera_heights=256, 38 | camera_widths=256, 39 | camera_depths=False, 40 | bin1_pos=(0.1, -0.26, 0.8), 41 | bin2_pos=(0.1, 0.13, 0.8), 42 | reward_scale=1.0, 43 | reward_shaping=False, 44 | ): 45 | 46 | self.subtask_id = 3 47 | super().__init__( 48 | robots, 49 | object_type, 50 | obstacle, 51 | bin1_pos, 52 | bin2_pos, 53 | env_configuration=env_configuration, 54 | controller_configs=controller_configs, 55 | mount_types=mount_types, 56 | gripper_types=gripper_types, 57 | initialization_noise=initialization_noise, 58 | use_camera_obs=use_camera_obs, 59 | use_object_obs=use_object_obs, 60 | use_task_id_obs=use_task_id_obs, 61 | has_renderer=has_renderer, 62 | has_offscreen_renderer=has_offscreen_renderer, 63 | render_camera=render_camera, 64 | render_collision_mesh=render_collision_mesh, 65 | render_visual_mesh=render_visual_mesh, 66 | render_gpu_device_id=render_gpu_device_id, 67 | control_freq=control_freq, 68 | horizon=horizon, 69 | ignore_done=ignore_done, 70 | hard_reset=hard_reset, 71 | camera_names=camera_names, 72 | camera_heights=camera_heights, 73 | camera_widths=camera_widths, 74 | camera_depths=camera_depths, 75 | reward_scale=reward_scale, 76 | reward_shaping=reward_shaping, 77 | ) 78 | 79 | def staged_rewards(self, action): 80 | """ 81 | Returns staged rewards based on current physical states. 82 | Stages consist of reaching, grasping, lifting, hovering and dropping. 83 | 84 | Returns: 85 | - (float) reaching reward 86 | - (float) grasping reward 87 | - (float) lifting reward 88 | - (float) hovering reward 89 | - (float) dropping reward 90 | """ 91 | 92 | reach_mult = 0.2 93 | grasp_mult = 0.3 94 | lift_mult = 0.5 95 | hover_mult = 0.7 96 | drop_mult = 0.95 97 | 98 | # reaching reward governed by distance to closest object 99 | r_reach = 0. 100 | if not self.object_in_trashcan: 101 | # get reaching reward via minimum distance to a target object 102 | dist = self._gripper_to_target( 103 | gripper=self.robots[0].gripper, 104 | target=self.object.root_body, 105 | target_type="body", 106 | return_distance=True, 107 | ) 108 | r_reach = (1 - np.tanh(10.0 * dist)) * reach_mult 109 | 110 | # grasping reward for touching any objects of interest 111 | r_grasp = 0 112 | if not self.object_in_trashcan: 113 | r_grasp = int(self._check_grasp( 114 | gripper=self.robots[0].gripper, 115 | object_geoms=[g for g in self.object.contact_geoms]) 116 | ) * grasp_mult 117 | 118 | # lifting reward for picking up an object 119 | r_lift = 0. 120 | if not self.object_in_trashcan and r_grasp > 0.: 121 | z_target = self.trashcan_pos[2] + \ 122 | self.trashcan_z_offset + self.trashcan_z_height + 0.2 123 | object_z_loc = self.sim.data.body_xpos[self.obj_body_id, 2] 124 | z_dist = np.abs(z_target - object_z_loc) 125 | r_lift = grasp_mult + (1 - np.tanh(5.0 * z_dist)) * ( 126 | lift_mult - grasp_mult 127 | ) 128 | 129 | object_x_loc = self.sim.data.body_xpos[self.obj_body_id, 0] 130 | object_y_loc = self.sim.data.body_xpos[self.obj_body_id, 1] 131 | object_z_loc = self.sim.data.body_xpos[self.obj_body_id, 2] 132 | 133 | trashcan_x_low = self.trashcan_pos[0] - self.trashcan_x_size 134 | trashcan_x_high = self.trashcan_pos[0] + self.trashcan_x_size 135 | x_check = trashcan_x_low < object_x_loc < trashcan_x_high 136 | 137 | trashcan_y_low = self.trashcan_pos[1] - self.trashcan_y_size 138 | trashcan_y_high = self.trashcan_pos[1] + self.trashcan_y_size 139 | y_check = trashcan_y_low < object_y_loc < trashcan_y_high 140 | 141 | # Not really sure what I'm doing here, so I'm copying from PickPlace 142 | dist = np.linalg.norm( 143 | np.array((self.trashcan_pos[0], self.trashcan_pos[1]) 144 | ) - np.array((object_x_loc, object_y_loc)) 145 | ) 146 | object_above_trashcan = x_check and y_check # TODO: add z check ? 147 | 148 | # hover reward for getting object in front of trashcan 149 | r_hover = 0. 150 | if not self.object_in_trashcan and r_lift > 0.45: 151 | if object_above_trashcan: 152 | r_hover = lift_mult + ( 153 | 1 - np.tanh(2.0 * dist) 154 | ) * (hover_mult - lift_mult) 155 | else: 156 | r_hover = r_lift + ( 157 | 1 - np.tanh(2.0 * dist) 158 | ) * (hover_mult - lift_mult) 159 | 160 | r_drop = 0 161 | if not self.object_in_trashcan and object_above_trashcan and action[-1] < 0: 162 | r_drop = 1 * drop_mult 163 | 164 | # porbably add a term to encourage opening the gripper 165 | 166 | return r_reach, r_grasp, r_lift, r_hover, r_drop 167 | 168 | def not_in_trashcan(self, obj_pos): 169 | """Checks whether object is inside the trashcan. 170 | 171 | Args: 172 | obj_pos (np.array): Current position of the object. 173 | 174 | Returns: 175 | bool: True if the object is NOT inside the trashcan. 176 | """ 177 | # get position of trashcan 178 | trashcan_x_low = self.trashcan_pos[0] - self.trashcan_x_size 179 | trashcan_x_high = self.trashcan_pos[0] + self.trashcan_x_size 180 | trashcan_y_low = self.trashcan_pos[1] - self.trashcan_y_size 181 | trashcan_y_high = self.trashcan_pos[1] + self.trashcan_y_size 182 | 183 | res = True 184 | if ( 185 | trashcan_x_low < obj_pos[0] < trashcan_x_high 186 | and trashcan_y_low < obj_pos[1] < trashcan_y_high 187 | and self.trashcan_z_offset + self.trashcan_pos[2] < obj_pos[2] < self.trashcan_z_offset + self.trashcan_pos[2] + 0.1 188 | ): 189 | res = False 190 | return res 191 | 192 | def _get_placement_initializer(self): 193 | """Helper function for defining placement initializer and object sampling bounds. 194 | """ 195 | super()._get_placement_initializer() 196 | 197 | # get limits of trashcan 198 | self.placement_initializer.append_sampler( 199 | sampler=UniformRandomSampler( 200 | name=f"{self.visual_object.name}ObjectSampler", 201 | mujoco_objects=self.visual_object, 202 | x_range=[0, 0], 203 | y_range=[0, 0], 204 | rotation=0., 205 | rotation_axis='z', 206 | ensure_object_boundary_in_range=False, 207 | ensure_valid_placement=False, 208 | reference_pos=self.trashcan_pos, 209 | z_offset=self.trashcan_z_offset, 210 | ) 211 | ) 212 | 213 | def _load_model(self): 214 | """Loads an xml model, puts it in self.model 215 | """ 216 | 217 | # load model for table top workspace 218 | self.mujoco_arena = TrashcanArena( 219 | ) 220 | 221 | # settings for trashcan position 222 | self.trashcan_pos = self.mujoco_arena.trashcan_pos 223 | self.trashcan_x_size = self.mujoco_arena.trashcan_floor_size[0] 224 | self.trashcan_y_size = self.mujoco_arena.trashcan_floor_size[1] 225 | self.trashcan_z_offset = self.mujoco_arena.trashcan_floor_size[2] 226 | self.trashcan_z_height = self.mujoco_arena.trashcan_height 227 | super()._load_model() 228 | 229 | # Generate placement initializer 230 | self._initialize_model() 231 | self._get_placement_initializer() 232 | 233 | def _setup_references(self): 234 | """Sets up references to important components. A reference is typically an 235 | index or a list of indices that point to the corresponding elements 236 | in a flatten array, which is how MuJoCo stores physical simulation data. 237 | """ 238 | self.object_in_trashcan = False 239 | super()._setup_references() 240 | 241 | def _reset_internal(self): 242 | """Resets simulation internal configurations. 243 | """ 244 | super()._reset_internal() 245 | 246 | def _check_success(self): 247 | """Check if object has been successfully placed in the trashcan. 248 | 249 | Returns: 250 | bool: True if object is placed correctly 251 | """ 252 | # remember if object is in the trashcan 253 | gripper_site_pos = self.sim.data.site_xpos[self.robots[0].eef_site_id] 254 | obj_pos = self.sim.data.body_xpos[self.obj_body_id] 255 | 256 | self.object_in_trashcan = not self.not_in_trashcan(obj_pos) 257 | gripper_not_in_bin = self.not_in_trashcan(gripper_site_pos) 258 | 259 | if gripper_not_in_bin: 260 | return self.object_in_trashcan 261 | else: 262 | return 0 263 | -------------------------------------------------------------------------------- /composuite/utils/demo_utils.py: -------------------------------------------------------------------------------- 1 | def choose_object(): 2 | """ 3 | Prints out object options, and returns the requested object. 4 | 5 | Returns: 6 | str: Requested object name 7 | """ 8 | # Get the list of objects 9 | objects = { 10 | "Box", 11 | "Dumbbell", 12 | "Plate", 13 | "Hollowbox" 14 | } 15 | 16 | # Make sure set is deterministically sorted 17 | objects = sorted(objects) 18 | 19 | # Select object 20 | print("Here is a list of available objects:\n") 21 | 22 | for k, obj in enumerate(objects): 23 | print("[{}] {}".format(k, obj)) 24 | print() 25 | try: 26 | s = input( 27 | "Choose an object " 28 | + "(enter a number from 0 to {}): ".format(len(objects) - 1) 29 | ) 30 | # parse input into a number within range 31 | k = min(max(int(s), 0), len(objects)) 32 | except: 33 | k = 0 34 | print("Input is not valid. Use {} by default.".format( 35 | list(objects)[k])) 36 | 37 | # Return requested object 38 | return list(objects)[k] 39 | 40 | 41 | def choose_obstacle(): 42 | """ 43 | Prints out obstacle options, and returns the requested obstacle. 44 | 45 | Returns: 46 | str: Requested obstacle name 47 | """ 48 | # Get the list of obstacles 49 | obstacles = { 50 | "None", 51 | "ObjectWall", 52 | "GoalWall", 53 | "ObjectDoor", 54 | } 55 | 56 | # Make sure set is deterministically sorted 57 | obstacles = sorted(obstacles) 58 | 59 | # Select obstacle 60 | print("Here is a list of available obstacles:\n") 61 | 62 | for k, obstacle in enumerate(obstacles): 63 | print("[{}] {}".format(k, obstacle)) 64 | print() 65 | try: 66 | s = input( 67 | "Choose an obstacle " 68 | + "(enter a number from 0 to {}): ".format(len(obstacles) - 1) 69 | ) 70 | # parse input into a number within range 71 | k = min(max(int(s), 0), len(obstacles)) 72 | except: 73 | k = 0 74 | print("Input is not valid. Use {} by default.".format( 75 | list(obstacles)[k])) 76 | 77 | # Return requested obstacle 78 | return list(obstacles)[k] 79 | 80 | 81 | def choose_task(): 82 | """ 83 | Prints out task options, and returns the requested task. 84 | 85 | Returns: 86 | str: Requested task name 87 | """ 88 | # Get the list of tasks 89 | tasks = { 90 | "PickPlace", 91 | "Shelf", 92 | "Push", 93 | "Trashcan" 94 | } 95 | 96 | # Make sure set is deterministically sorted 97 | tasks = sorted(tasks) 98 | 99 | # Select task 100 | print("Here is a list of available tasks:\n") 101 | 102 | for k, task in enumerate(tasks): 103 | print("[{}] {}".format(k, task)) 104 | print() 105 | try: 106 | s = input( 107 | "Choose an task " 108 | + "(enter a number from 0 to {}): ".format(len(tasks) - 1) 109 | ) 110 | # parse input into a number within range 111 | k = min(max(int(s), 0), len(tasks)) 112 | except: 113 | k = 0 114 | print("Input is not valid. Use {} by default.".format(list(tasks)[k])) 115 | 116 | # Return requested task 117 | return list(tasks)[k] 118 | 119 | def choose_robots(exclude_bimanual=False): 120 | """ 121 | Prints out robot options, and returns the requested robot. Restricts options to single-armed robots if 122 | @exclude_bimanual is set to True (False by default) 123 | 124 | Args: 125 | exclude_bimanual (bool): If set, excludes bimanual robots from the robot options 126 | 127 | Returns: 128 | str: Requested robot name 129 | """ 130 | # Get the list of robots 131 | robots = { 132 | "Panda", 133 | "Jaco", 134 | "Kinova3", 135 | "IIWA", 136 | } 137 | 138 | # Add Baxter if bimanual robots are not excluded 139 | if not exclude_bimanual: 140 | robots.add("Baxter") 141 | 142 | # Make sure set is deterministically sorted 143 | robots = sorted(robots) 144 | 145 | # Select robot 146 | print("Here is a list of available robots:\n") 147 | 148 | for k, robot in enumerate(robots): 149 | print("[{}] {}".format(k, robot)) 150 | print() 151 | try: 152 | s = input( 153 | "Choose a robot " 154 | + "(enter a number from 0 to {}): ".format(len(robots) - 1) 155 | ) 156 | # parse input into a number within range 157 | k = min(max(int(s), 0), len(robots)) 158 | except: 159 | k = 0 160 | print("Input is not valid. Use {} by default.".format(list(robots)[k])) 161 | 162 | # Return requested robot 163 | return list(robots)[k] -------------------------------------------------------------------------------- /composuite/utils/mjk_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import composuite 3 | 4 | def xml_path_completion(xml_path): 5 | """ 6 | Takes in a local xml path and returns a full path. 7 | if @xml_path is absolute, do nothing 8 | if @xml_path is not absolute, load xml that is shipped by the package 9 | 10 | Args: 11 | xml_path (str): local xml path 12 | 13 | Returns: 14 | str: Full (absolute) xml path 15 | """ 16 | if xml_path.startswith("/"): 17 | full_path = xml_path 18 | else: 19 | full_path = os.path.join(composuite.assets_root, xml_path) 20 | return full_path -------------------------------------------------------------------------------- /example_scripts/train_ppo.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import argparse 3 | import os 4 | import json 5 | import warnings 6 | 7 | import torch 8 | 9 | import composuite 10 | 11 | from spinup.algos.pytorch.ppo.core import MLPActorCritic 12 | from spinup.algos.pytorch.ppo.compositional_core import CompositionalMLPActorCritic 13 | from spinup.algos.pytorch.ppo.ppo import ppo 14 | from spinup.utils.run_utils import setup_logger_kwargs 15 | from spinup.utils.mpi_tools import proc_id 16 | 17 | 18 | def parse_default_args(): 19 | parser = argparse.ArgumentParser() 20 | 21 | # Directory information 22 | parser.add_argument('--data-dir', default='spinningup_training') 23 | parser.add_argument('--load-dir', default=None) 24 | 25 | # Neural network training parameters 26 | parser.add_argument('--exp-name', type=str, default='ppo') 27 | parser.add_argument('--seed', '-s', type=int, default=42) 28 | parser.add_argument('--gamma', type=float, default=0.99) 29 | parser.add_argument('--steps', type=int, default=16000) 30 | parser.add_argument('--epochs', type=int, default=625) 31 | parser.add_argument('--clip', type=float, default=0.2) 32 | parser.add_argument('--pi-lr', type=float, default=1e-4) 33 | parser.add_argument('--vf-lr', type=float, default=1e-4) 34 | parser.add_argument('--pi-iters', type=int, default=128) 35 | parser.add_argument('--vf-iters', type=int, default=128) 36 | parser.add_argument('--target-kl', type=float, default=0.02) 37 | parser.add_argument('--ent-coef', type=float, default=0.) 38 | parser.add_argument('--log-std-init', type=float, default=0.) 39 | 40 | # Architecture information 41 | parser.add_argument('--hid', type=int, default=64) 42 | parser.add_argument('--l', type=int, default=2) 43 | 44 | # Task information 45 | parser.add_argument('--controller', type=str, default="joint") 46 | 47 | parser.add_argument('--robot', type=str, default="IIWA", 48 | choices=['IIWA', 'Jaco', 'Kinova3', 'Panda'], 49 | help='Robotic manipulator to use in task.') 50 | parser.add_argument('--object', type=str, default="Hollowbox", 51 | choices=['Hollowbox', 'Box', 'Dumbbell', 'Plate'], 52 | help='Object to use in the environment.') 53 | parser.add_argument('--obstacle', type=str, default='None', 54 | choices=['None', 'ObjectDoor', 55 | 'ObjectWall', 'Goalwall'], 56 | help='Obstacle to use in environment.') 57 | parser.add_argument('--task', type=str, default="PickPlace", 58 | choices=['PickPlace', 'Push', 'Shelf', 'Trashcan'], 59 | help='Objective to train on.') 60 | parser.add_argument('--horizon', type=int, default=500, 61 | help='Number of steps before the environment resets') 62 | parser.add_argument('--task-id', type=int, default=-1, 63 | help='This ID is used for the single task \ 64 | learner as an unravel index \ 65 | into the possible task configurations. \ 66 | It can take a value from 0 to 255. The \ 67 | default of -1 means that it is not used. \ 68 | If this is used, the args for robot, \ 69 | object, task and obstacle will be \ 70 | overwritten by the specific configuration. \ 71 | This main purpose of the flag is to run all \ 72 | possible single task experiments by running \ 73 | all indices from 0 to 255.') 74 | parser.add_argument('--num-train-tasks', type=int, default=1) 75 | 76 | # Learner type 77 | parser.add_argument('--learner-type', type=str, default='stl', 78 | choices=['stl', 'mtl', 'comp'], 79 | help='stl: Single-task learner, \ 80 | mtl: Multi-task learner, \ 81 | comp: Compositional learner') 82 | 83 | # Experiment type 84 | parser.add_argument('--experiment-type', default='default', 85 | choices=['default', 'smallscale', 'holdout']) 86 | parser.add_argument('--smallscale-elem', default='IIWA') 87 | parser.add_argument('--holdout-elem', default='IIWA') 88 | 89 | args = parser.parse_args() 90 | 91 | if args.experiment_type == 'default': 92 | args.exp_name = os.path.join( 93 | args.exp_name, args.experiment_type + 94 | "_" + str(args.num_train_tasks), 95 | args.learner_type, 's' + str(args.seed)) 96 | elif args.experiment_type == 'holdout': 97 | args.exp_name = os.path.join( 98 | args.exp_name, args.experiment_type + 99 | "_" + str(args.num_train_tasks), 100 | args.learner_type, args.holdout_elem, 's' + str(args.seed)) 101 | elif args.experiment_type == 'smallscale': 102 | args.exp_name = os.path.join( 103 | args.exp_name, args.experiment_type + 104 | "_" + str(args.num_train_tasks), 105 | args.learner_type, args.smallscale_elem, 's' + str(args.seed)) 106 | 107 | assert args.task_id >= -1 and args.task_id <= 255, "Task ID out of bounds" 108 | if args.learner_type == 'stl': 109 | assert args.experiment_type == 'default', \ 110 | "You have selected the single-task agent but are \ 111 | trying to run a multi-task experiment." 112 | 113 | if args.task_id != -1 and args.learner_type != 'stl': 114 | warnings.warn("You have selected a specific task ID while using \ 115 | a learner that is not the single-task learner. \ 116 | Ignoring task ID.") 117 | return args 118 | 119 | 120 | def train_model(robot, obj, obstacle, task, args, logger_kwargs): 121 | """Run a single model training. 122 | 123 | Args: 124 | robot (str): Robot to use in environment. 125 | obj (str): Object to use in environment. 126 | obstacle (str): Obstacle to use in environment. 127 | task (str): Objective to use in environment. 128 | args (Namespace): Training namespace arguments. 129 | logger_kwargs (dict): Additional kwargs for spinningup ppo. 130 | 131 | Raises: 132 | NotImplementedError: Raises if unknown 133 | learner type is selected. 134 | """ 135 | torch.backends.cudnn.deterministic = True 136 | torch.backends.cudnn.benchmark = False 137 | torch.set_num_threads(1) 138 | 139 | checkpoint = None 140 | if args.load_dir is not None: 141 | checkpoint = torch.load(os.path.join( 142 | args.load_dir, 'pyt_save', 'state_dicts.pt')) 143 | 144 | def env_fn(): 145 | return composuite.make( 146 | robot, obj, obstacle, task, 147 | args.controller, args.horizon, use_task_id_obs=True) 148 | 149 | if args.learner_type == 'comp': 150 | network_class = CompositionalMLPActorCritic 151 | hidden_sizes = ((32,), (32, 32), (64, 64, 64), (64, 64, 64)) 152 | 153 | ac_kwargs = { 154 | 'log_std_init': args.log_std_init, 155 | 'hidden_sizes': hidden_sizes, 156 | 'module_names': ['obstacle_id', 'object_id', 157 | 'subtask_id', 'robot_id'], 158 | 'module_input_names': ['obstacle-state', 'object-state', 159 | 'goal-state', 'robot0_proprio-state'], 160 | 'interface_depths': [-1, 1, 2, 3], 161 | 'graph_structure': [[0], [1], [2], [3]], 162 | } 163 | elif args.learner_type == 'stl' or args.learner_type == 'mtl': 164 | network_class = MLPActorCritic 165 | ac_kwargs = dict( 166 | hidden_sizes=[args.hid]*args.l, log_std_init=args.log_std_init) 167 | else: 168 | raise NotImplementedError("Unknown learner type was selected. \ 169 | Options are stl, mtl and comp.") 170 | 171 | ppo(env_fn=env_fn, 172 | actor_critic=network_class, 173 | ac_kwargs=ac_kwargs, 174 | seed=args.seed, gamma=args.gamma, steps_per_epoch=args.steps, 175 | epochs=args.epochs, clip_ratio=args.clip, 176 | pi_lr=args.pi_lr, vf_lr=args.vf_lr, 177 | train_pi_iters=args.pi_iters, train_v_iters=args.vf_iters, 178 | target_kl=args.target_kl, logger_kwargs=logger_kwargs, 179 | max_ep_len=args.horizon, ent_coef=args.ent_coef, 180 | log_per_proc=args.learner_type != 'stl', checkpoint=checkpoint) 181 | 182 | 183 | def main(): 184 | """Runs the training script. 185 | 186 | Raises: 187 | NotImplementedError: Raises if selected 188 | experiment type does not exist. 189 | """ 190 | args = parse_default_args() 191 | np.random.seed(args.seed) 192 | if args.experiment_type == 'default': 193 | if args.learner_type == 'stl': 194 | if args.task_id != -1: 195 | train_tasks, test_tasks = composuite.sample_tasks( 196 | experiment_type=args.experiment_type, 197 | num_train=0, 198 | no_shuffle=True 199 | ) 200 | train_task = test_tasks[args.task_id] 201 | robot, obj, obstacle, task = train_task 202 | test_tasks.remove(train_task) 203 | train_tasks.append(train_task) 204 | else: 205 | robot = args.robot 206 | obj = args.object 207 | task = args.task 208 | obstacle = args.obstacle 209 | 210 | train_tasks = [(robot, obj, obstacle, task)] 211 | _, test_tasks = composuite.sample_tasks( 212 | experiment_type=args.experiment_type, 213 | num_train=0, 214 | no_shuffle=True 215 | ) 216 | test_tasks.remove(train_tasks[0]) 217 | else: 218 | train_tasks, test_tasks = composuite.sample_tasks( 219 | experiment_type=args.experiment_type, 220 | num_train=args.num_train_tasks, 221 | shuffling_seed=args.seed, 222 | no_shuffle=False 223 | ) 224 | robot, obj, obstacle, task = train_tasks[proc_id()] 225 | elif args.experiment_type == 'smallscale' or \ 226 | args.experiment_type == 'holdout': 227 | train_tasks, test_tasks = composuite.sample_tasks( 228 | experiment_type=args.experiment_type, 229 | num_train=args.num_train_tasks, 230 | smallscale_elem=args.smallscale_elem, 231 | holdout_elem=args.holdout_elem, 232 | shuffling_seed=args.seed, 233 | no_shuffle=False 234 | ) 235 | robot, obj, obstacle, task = train_tasks[proc_id()] 236 | else: 237 | raise NotImplementedError("Specified experiment type does not exist.") 238 | 239 | args.robot = robot 240 | args.object = obj 241 | args.obstacle = obstacle 242 | args.task = task 243 | 244 | os.makedirs(os.path.join(args.data_dir, args.exp_name), exist_ok=True) 245 | with open( 246 | os.path.join( 247 | args.data_dir, args.exp_name, 'args_{}.json'.format(proc_id())), 248 | 'w') as f: 249 | json.dump(args.__dict__, f, indent=2) 250 | 251 | if proc_id() == 0: 252 | with open(os.path.join(args.data_dir, args.exp_name, 'tasks.json'), 'w') as f: 253 | tasks = { 254 | 'train_tasks': train_tasks, 255 | 'test_tasks': test_tasks 256 | } 257 | json.dump(tasks, f, indent=2) 258 | 259 | logger_kwargs = setup_logger_kwargs( 260 | args.exp_name, data_dir=args.data_dir) 261 | 262 | train_model( 263 | robot=robot, obj=obj, obstacle=obstacle, task=task, 264 | args=args, logger_kwargs=logger_kwargs) 265 | 266 | 267 | if __name__ == '__main__': 268 | main() 269 | -------------------------------------------------------------------------------- /requirements_default.txt: -------------------------------------------------------------------------------- 1 | robosuite==1.4.0 2 | termcolor==2.4.0 3 | -------------------------------------------------------------------------------- /requirements_paper.txt: -------------------------------------------------------------------------------- 1 | git+https://github.com/ARISE-Initiative/robosuite@7a92b7071eb4922f26fe75b6049c05ca967cb175#egg=robosuite 2 | mujoco-py==2.0.2.13 3 | -e . 4 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from os import path 2 | from setuptools import find_packages, setup 3 | 4 | # read the contents of README file 5 | this_directory = path.abspath(path.dirname(__file__)) 6 | with open(path.join(this_directory, "README.md"), encoding="utf-8") as f: 7 | lines = f.readlines() 8 | 9 | # remove images from README 10 | lines = [x for x in lines if ".png" not in x] 11 | long_description = "".join(lines) 12 | 13 | setup( 14 | name="composuite", 15 | packages=[package for package in find_packages() if package.startswith("composuite")], 16 | install_requires=[ 17 | "robosuite", 18 | "gym>=0.15.7", 19 | "h5py>=2.10.0", 20 | ], 21 | eager_resources=["*"], 22 | include_package_data=True, 23 | python_requires=">=3,<3.12", 24 | description="CompoSuite: A Compositional Reinforcement Learning Benchmark", 25 | author="Jorge Mendez, Marcel Hussing, Meghna Gummadi and Eric Eaton", 26 | url="https://github.com/Lifelong-ML/CompoSuite", 27 | author_email="", 28 | version="1.0", 29 | long_description=long_description, 30 | long_description_content_type="text/markdown", 31 | ) 32 | -------------------------------------------------------------------------------- /tests/test_main_functions.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from composuite import sample_tasks 3 | 4 | AVAILABLE_ROBOTS = ["IIWA", "Jaco", "Kinova3", "Panda"] 5 | AVAILABLE_OBSTACLES = ["None", "GoalWall", 6 | "ObjectDoor", "ObjectWall"] 7 | AVAILABLE_OBJECTS = ["Box", "Dumbbell", "Plate", "Hollowbox"] 8 | AVAILABLE_TASKS = ["PickPlace", "Push", "Shelf", "Trashcan"] 9 | 10 | 11 | def test_sample_tasks_smallscale(): 12 | possible_elements = AVAILABLE_ROBOTS + AVAILABLE_OBSTACLES + \ 13 | AVAILABLE_OBJECTS + AVAILABLE_TASKS 14 | for e in possible_elements: 15 | train_tasks, test_tasks = sample_tasks( 16 | experiment_type='smallscale', num_train=32, smallscale_elem=e 17 | ) 18 | 19 | for t in train_tasks: 20 | assert e in t 21 | for t in test_tasks: 22 | assert e in t 23 | 24 | 25 | def test_sample_tasks_holdout(): 26 | possible_elements = AVAILABLE_ROBOTS + AVAILABLE_OBSTACLES + \ 27 | AVAILABLE_OBJECTS + AVAILABLE_TASKS 28 | for e in possible_elements: 29 | train_tasks, test_tasks = sample_tasks( 30 | experiment_type='holdout', num_train=56, holdout_elem=e 31 | ) 32 | assert len(train_tasks) == 56 33 | assert len(test_tasks) == 63 34 | 35 | num_train_tasks_with_e = 0 36 | for t in train_tasks: 37 | if e in t: 38 | num_train_tasks_with_e += 1 39 | assert num_train_tasks_with_e == 1 40 | 41 | for t in test_tasks: 42 | assert e in t 43 | 44 | 45 | def test_sample_tasks_num_train(): 46 | num_train = 0 47 | train_tasks, test_tasks = sample_tasks(num_train=num_train) 48 | assert len(train_tasks) == num_train 49 | assert len(test_tasks) == 256 - num_train 50 | 51 | num_train = 5 52 | train_tasks, test_tasks = sample_tasks(num_train=num_train) 53 | assert len(train_tasks) == num_train 54 | assert len(test_tasks) == 256 - num_train 55 | 56 | num_train = 256 57 | train_tasks, test_tasks = sample_tasks(num_train=num_train) 58 | assert len(train_tasks) == num_train 59 | assert len(test_tasks) == 256 - num_train 60 | 61 | with pytest.raises(AssertionError): 62 | num_train = -1 63 | train_tasks, test_tasks = sample_tasks(num_train=num_train) 64 | 65 | with pytest.raises(AssertionError): 66 | num_train = 257 67 | train_tasks, test_tasks = sample_tasks(num_train=num_train) 68 | 69 | 70 | def test_sample_tasks_num_train_smallscale_1(): 71 | possible_elements = AVAILABLE_ROBOTS + AVAILABLE_OBSTACLES + \ 72 | AVAILABLE_OBJECTS + AVAILABLE_TASKS 73 | num_train = 0 74 | for e in possible_elements: 75 | train_tasks, test_tasks = sample_tasks( 76 | experiment_type='smallscale', num_train=num_train, smallscale_elem=e 77 | ) 78 | assert len(train_tasks) == num_train 79 | assert len(test_tasks) == 64 - num_train 80 | 81 | 82 | def test_sample_tasks_num_train_smallscale_2(): 83 | possible_elements = AVAILABLE_ROBOTS + AVAILABLE_OBSTACLES + \ 84 | AVAILABLE_OBJECTS + AVAILABLE_TASKS 85 | num_train = -1 86 | for e in possible_elements: 87 | with pytest.raises(AssertionError): 88 | train_tasks, test_tasks = sample_tasks( 89 | experiment_type='smallscale', num_train=num_train, smallscale_elem=e 90 | ) 91 | 92 | num_train = 65 93 | for e in possible_elements: 94 | with pytest.raises(AssertionError): 95 | train_tasks, test_tasks = sample_tasks( 96 | experiment_type='smallscale', num_train=num_train, smallscale_elem=e 97 | ) 98 | 99 | 100 | def test_sample_tasks_num_train_holdout_1(): 101 | possible_elements = AVAILABLE_ROBOTS + AVAILABLE_OBSTACLES + \ 102 | AVAILABLE_OBJECTS + AVAILABLE_TASKS 103 | num_train = 0 104 | for e in possible_elements: 105 | with pytest.raises(AssertionError): 106 | train_tasks, test_tasks = sample_tasks( 107 | experiment_type='holdout', num_train=num_train, holdout_elem=e 108 | ) 109 | --------------------------------------------------------------------------------