[Docs] New documentation on scene datasets

docs/source/user_guide/datasets/scenes.md

@@ -2,10 +2,38 @@
 
 We provide a command line tool to download scene datasets (typically adapted from the original dataset). 
 
-We can support loading just about any scene dataset but currently only provide an option to load ReplicaCAD as it is the best tested one (with interactable objects) we have access to. We are in the process still of processing AI2THOR scene datasets to make them available as well in ManiSkill and standardizing a Scene-level task / definition API to make it easier to support running GPU parallelized simulation in any scene
+ManiSkill can build any scene provided assets are provided. ManiSkill out of the box provides code and download links to use the [ReplicaCAD](https://aihabitat.org/datasets/replica_cad/) and [AI2THOR](https://github.com/allenai/ai2thor) set of scenes (shown below). These are picked because we are able to make them *interactive* scenes where objects can be manipulated and moved around.
+
+```{figure} images/two_scenes_examples.png
+```
 
 ```bash
 # list all scene datasets available for download
 python -m mani_skill.utils.download_asset --list "scene"
-python -m mani_skill.utils.download_asset ReplicaCAD
+python -m mani_skill.utils.download_asset ReplicaCAD # small scene and fast to download
+python -m mani_skill.utils.download_asset AI2THOR # lots of scenes and slow to download
 ```
+
+## Exploring the Scene Datasets
+
+To explore the scene datasets, you can provide an environment ID and a seed (to change which scene is sampled if there are several available) and run the random action script. Shown below are the two environment IDs configured already to enable you to play with ReplicaCAD and ArchitecTHOR, one of the scene sets in AI2THOR.
+
+```bash
+python -m mani_skill.examples.demo_random_action \
+  -e "ReplicaCAD_SceneManipulation-v1" \
+  --render-mode="rgb_array" --record-dir="videos" # run headless and save video
+
+python -m mani_skill.examples.demo_random_action \
+  -e "ArchitecTHOR_SceneManipulation-v1" --render-mode="human" \
+  -s 3 # open a GUI and sample a scene with seed 3
+```
+
+## Training on the Scene Datasets
+
+Large scene datasets with hundreds of objects like ReplicaCAD and AI2THOR can be used to train more general purpose robots/agents and also serve as synthetic data generation sources. We are still in the process of providing more example code and documentation about how to best leverage these scene datasets but for now we provide code to explore and interact with the scene datasets.
+
+### Reinforcement Learning / Imitation Learning
+
+We are currently in the process of building task code similar to the ReplicaCAD Rearrange challenge and will open source that when it is complete. Otherwise at the moment there are not any trainable tasks with defined success/fail conditions and/or rewards that use any of the big scene datasets.
+
+### Computer Vision / Synthetic 2D/3D Data Generation (WIP)

mani_skill/envs/scenes/base_env.py

@@ -1,7 +1,6 @@
-from typing import Any, Dict, Union, List
+from typing import Any, Dict, Union
 
 import numpy as np
-import torch
 import sapien as sapien
 import sapien.physx as physx
 import torch
@@ -20,14 +19,15 @@
 @register_env("SceneManipulation-v1", max_episode_steps=200)
 class SceneManipulationEnv(BaseEnv):
     """
-    A base environment for simulating manipulation tasks in more complex scenes. Creating this base environment is only useful for explorations/visualization, there are no success/failure
-    metrics or rewards.
+    A base environment for simulating manipulation tasks in more complex scenes. Creating this base environment is only useful
+    for explorations/visualization, there are no success/failure metrics or rewards.
 
     Args:
         robot_uids: Which robot to place into the scene. Default is "fetch"
 
-        fixed_scene: whether to build static set of scenes.
-        Default is True as reconfiguring is expensive. When fixed_scene=True, one can rebuild scenes via env.reset(seed=seed, options=dict(reconfigure=True))
+        fixed_scene:
+            When True, will never reconfigure the environment during resets unless you run env.reset(seed=seed, options=dict(reconfigure=True))
+            and explicitly reconfigure. If False, will reconfigure every reset.
 
         scene_builder_cls: Scene builder class to build a scene with. Default is the ArchitecTHORSceneBuilder which builds a scene from AI2THOR.
             Any of the AI2THOR SceneBuilders are supported in this environment
@@ -43,21 +43,32 @@ def __init__(
         self,
         *args,
         robot_uids="fetch",
-        fixed_scene=True,
         scene_builder_cls: Union[str, SceneBuilder] = "ReplicaCAD",
         build_config_idxs=None,
         init_config_idxs=None,
+        num_envs=1,
+        reconfiguration_freq=None,
         **kwargs
     ):
-        self.fixed_scene = fixed_scene
         if isinstance(scene_builder_cls, str):
             scene_builder_cls = REGISTERED_SCENE_BUILDERS[
                 scene_builder_cls
             ].scene_builder_cls
         self.scene_builder: SceneBuilder = scene_builder_cls(self)
         self.build_config_idxs = build_config_idxs
         self.init_config_idxs = init_config_idxs
-        super().__init__(*args, robot_uids=robot_uids, **kwargs)
+        if reconfiguration_freq is None:
+            if num_envs == 1:
+                reconfiguration_freq = 1
+            else:
+                reconfiguration_freq = 0
+        super().__init__(
+            *args,
+            robot_uids=robot_uids,
+            reconfiguration_freq=reconfiguration_freq,
+            num_envs=num_envs,
+            **kwargs
+        )
 
     @property
     def _default_sim_config(self):
@@ -74,10 +85,6 @@ def reset(self, seed=None, options=None):
         self._set_episode_rng(seed)
         if options is None:
             options = dict(reconfigure=False)
-        if "reconfigure" not in options:
-            options["reconfigure"] = False
-        if not self.fixed_scene:
-            options["reconfigure"] = True
         if "reconfigure" in options and options["reconfigure"]:
             self.build_config_idxs = options.pop(
                 "build_config_idxs", self.build_config_idxs
@@ -130,10 +137,10 @@ def compute_normalized_dense_reward(
     @property
     def _default_sensor_configs(self):
         if self.robot_uids == "fetch":
-            return ()
+            return []
 
         pose = sapien_utils.look_at([0.3, 0, 0.6], [-0.1, 0, 0.1])
-        return CameraConfig("base_camera", pose, 128, 128, np.pi / 2, 0.01, 100)
+        return [CameraConfig("base_camera", pose, 128, 128, np.pi / 2, 0.01, 100)]
 
     @property
     def _default_human_render_camera_configs(self):

mani_skill/utils/download_asset.py

@@ -157,6 +157,12 @@ def initialize_sources():
         target_path="scene_datasets/replica_cad_dataset/rearrange",
     )
 
+    DATA_SOURCES["AI2THOR"] = DataSource(
+        source_type="scene",
+        url="https://huggingface.co/datasets/haosulab/AI2THOR/resolve/main/ai2thor.zip",
+        target_path="scene_datasets/ai2thor",
+    )
+
 
 def initialize_extra_sources():
     DATA_SOURCES["xmate3_robotiq"] = DataSource(

mani_skill/utils/scene_builder/ai2thor/scene_builder.py

@@ -8,19 +8,19 @@
 from pathlib import Path
 from typing import Dict, List, Tuple, Union
 
-import torch
 import numpy as np
 import sapien
 import sapien.core as sapien
 import sapien.physx as physx
+import torch
 import transforms3d
 from tqdm import tqdm
 
 from mani_skill import ASSET_DIR
 from mani_skill.agents.robots.fetch import (
-    Fetch,
-    FETCH_WHEELS_COLLISION_BIT,
     FETCH_BASE_COLLISION_BIT,
+    FETCH_WHEELS_COLLISION_BIT,
+    Fetch,
 )
 from mani_skill.envs.scene import ManiSkillScene
 from mani_skill.utils.scene_builder import SceneBuilder
@@ -40,9 +40,16 @@
     None,
 )
 
-# TODO (arth): fix coacd so this isn't necessary
-WORKING_OBJS = ["apple", "potato", "tomato"]
-
+WORKING_OBJS = [
+    "apple",
+    "potato",
+    "tomato",
+    "lettuce",
+    "soap",
+    "sponge",
+    "plate",
+    "book",
+]
 FETCH_BUILD_CONFIG_IDX_TO_START_POS = {
     0: (-3, 0),
     1: (-2, -2),
@@ -233,6 +240,8 @@ def build(
                 if npy_fp.exists():
                     self._navigable_positions[bci] = np.load(npy_fp)
 
+        self.scene.set_ambient_light([0.3, 0.3, 0.3])
+
         # merge actors into one
         self.bg = Actor.create_from_entities(
             bgs,

mani_skill/utils/scene_builder/ai2thor/variants.py

@@ -1,10 +1,13 @@
+from mani_skill.utils.scene_builder.registration import register_scene_builder
+
 from .scene_builder import AI2THORBaseSceneBuilder
 
 
 class ProcTHORSceneBuilder(AI2THORBaseSceneBuilder):
     scene_dataset = "ProcTHOR"
 
 
+@register_scene_builder("ArchitecTHOR")
 class ArchitecTHORSceneBuilder(AI2THORBaseSceneBuilder):
     scene_dataset = "ArchitecTHOR"