Merge pull request #328 from StanfordVL/feat/external-sensors

Feat/external sensors

omnigibson/configs/default_cfg.yaml

@@ -1,4 +1,7 @@
-env: {}
+env:
+  action_timestep: 0.01666              # (float): environment executes action per action_timestep second
+  physics_timestep: 0.01666             # (float): physics timestep for physx
+  device: null                          # (None or str): specifies the device to be used if running on the gpu with torch backend
 
 render:
   viewer_width: 1280

omnigibson/configs/fetch_behavior.yaml

@@ -1,5 +1,21 @@
 env:
+  action_timestep: 0.01666              # (float): environment executes action per action_timestep second
+  physics_timestep: 0.01666             # (float): physics timestep for physx
+  device: null                          # (None or str): specifies the device to be used if running on the gpu with torch backend
+  automatic_reset: false                # (bool): whether to automatic reset after an episode finishes
+  flatten_action_space: false           # (bool): whether to flatten the action space as a sinle 1D-array
+  flatten_obs_space: false              # (bool): whether the observation space should be flattened when generated
+  use_external_obs: false               # (bool): Whether to use external observations or not
   initial_pos_z_offset: 0.1
+  external_sensors:                     # (None or list): If specified, list of sensor configurations for external sensors to add. Should specify sensor "type" and any additional kwargs to instantiate the sensor. Each entry should be the kwargs passed to @create_sensor, in addition to position, orientation
+    - sensor_type: VisionSensor
+      modalities: [rgb, depth]
+      sensor_kwargs:
+        image_height: 128
+        image_width: 128
+      local_position: [0, 0, 1.0]
+      local_orientation: [0.707, 0.0, 0.0, 0.707]
+
 
 render:
   viewer_width: 1280

omnigibson/configs/tiago_primitives.yaml

@@ -1,5 +1,13 @@
 env:
+  action_timestep: 0.01666              # (float): environment executes action per action_timestep second
+  physics_timestep: 0.01666             # (float): physics timestep for physx
+  device: null                          # (None or str): specifies the device to be used if running on the gpu with torch backend
+  automatic_reset: false                # (bool): whether to automatic reset after an episode finishes
+  flatten_action_space: false           # (bool): whether to flatten the action space as a sinle 1D-array
+  flatten_obs_space: false              # (bool): whether the observation space should be flattened when generated
+  use_external_obs: false               # (bool): Whether to use external observations or not
   initial_pos_z_offset: 0.1
+  external_sensors: null                # (None or list): If specified, list of sensor configurations for external sensors to add. Should specify sensor "type" and any additional kwargs to instantiate the sensor. Each entry should be the kwargs passed to @create_sensor, in addition to position, orientation
 
 render:
   viewer_width: 1280

omnigibson/configs/turtlebot_nav.yaml

@@ -1,5 +1,13 @@
 env:
+  action_timestep: 0.01666              # (float): environment executes action per action_timestep second
+  physics_timestep: 0.01666             # (float): physics timestep for physx
+  device: null                          # (None or str): specifies the device to be used if running on the gpu with torch backend
+  automatic_reset: false                # (bool): whether to automatic reset after an episode finishes
+  flatten_action_space: false           # (bool): whether to flatten the action space as a sinle 1D-array
+  flatten_obs_space: false              # (bool): whether the observation space should be flattened when generated
+  use_external_obs: false               # (bool): Whether to use external observations or not
   initial_pos_z_offset: 0.1
+  external_sensors: null                # (None or list): If specified, list of sensor configurations for external sensors to add. Should specify sensor "type" and any additional kwargs to instantiate the sensor. Each entry should be the kwargs passed to @create_sensor, in addition to position, orientation
 
 render:
   viewer_width: 1280

omnigibson/envs/env_base.py

@@ -8,6 +8,7 @@
 from omnigibson.scene_graphs.graph_builder import SceneGraphBuilder
 from omnigibson.tasks import REGISTERED_TASKS
 from omnigibson.scenes import REGISTERED_SCENES
+from omnigibson.sensors import create_sensor
 from omnigibson.utils.gym_utils import GymObservable, recursively_generate_flat_dict
 from omnigibson.utils.config_utils import parse_config
 from omnigibson.utils.ui_utils import create_module_logger
@@ -23,42 +24,20 @@ class Environment(gym.Env, GymObservable, Recreatable):
     """
     Core environment class that handles loading scene, robot(s), and task, following OpenAI Gym interface.
     """
-    def __init__(
-        self,
-        configs,
-        action_timestep=1 / 60.0,
-        physics_timestep=1 / 60.0,
-        device=None,
-        automatic_reset=False,
-        flatten_action_space=False,
-        flatten_obs_space=False,
-    ):
+    def __init__(self, configs):
         """
         Args:
             configs (str or dict or list of str or dict): config_file path(s) or raw config dictionaries.
                 If multiple configs are specified, they will be merged sequentially in the order specified.
-                This allows procedural generation of a "full" config from small sub-configs.
-            action_timestep (float): environment executes action per action_timestep second
-            physics_timestep: physics timestep for physx
-            device (None or str): specifies the device to be used if running on the gpu with torch backend
-            automatic_reset (bool): whether to automatic reset after an episode finishes
-            flatten_action_space (bool): whether to flatten the action space as a sinle 1D-array
-            flatten_obs_space (bool): whether the observation space should be flattened when generated
+                This allows procedural generation of a "full" config from small sub-configs. For valid keys, please
+                see @default_config below
         """
         # Call super first
         super().__init__()
 
-        # Store settings and other initialized values
-        self._automatic_reset = automatic_reset
-        self._flatten_action_space = flatten_action_space
-        self._flatten_obs_space = flatten_obs_space
-        self.physics_timestep = physics_timestep
-        self.action_timestep = action_timestep
-        self.device = device
-
         # Initialize other placeholders that will be filled in later
-        self._initial_pos_z_offset = None                   # how high to offset object placement to account for one action step of dropping
         self._task = None
+        self._external_sensors = None
         self._loaded = None
         self._current_episode = 0
 
@@ -75,6 +54,15 @@ def __init__(
         for config in configs:
             merge_nested_dicts(base_dict=self.config, extra_dict=parse_config(config), inplace=True)
 
+        # Store settings and other initialized values
+        self._automatic_reset = self.env_config["automatic_reset"]
+        self._flatten_action_space = self.env_config["flatten_action_space"]
+        self._flatten_obs_space = self.env_config["flatten_obs_space"]
+        self.physics_timestep = self.env_config["physics_timestep"]
+        self.action_timestep = self.env_config["action_timestep"]
+        self.device = self.env_config["device"]
+        self._initial_pos_z_offset = self.env_config["initial_pos_z_offset"]    # how high to offset object placement to account for one action step of dropping
+
         # Create the scene graph builder
         self._scene_graph_builder = None
         if "scene_graph" in self.config and self.config["scene_graph"] is not None:
@@ -280,6 +268,33 @@ def _load_objects(self):
 
         assert og.sim.is_stopped(), "Simulator must be stopped after loading objects!"
 
+    def _load_external_sensors(self):
+        """
+        Load any additional custom external sensors into the scene
+        """
+        assert og.sim.is_stopped(), "Simulator must be stopped before loading external sensors!"
+        sensors_config = self.env_config["external_sensors"]
+        if sensors_config is not None:
+            self._external_sensors = dict()
+            for i, sensor_config in enumerate(sensors_config):
+                # Add a name for the object if necessary
+                if "name" not in sensor_config:
+                    sensor_config["name"] = f"external_sensor{i}"
+                # Determine prim path if not specified
+                if "prim_path" not in sensor_config:
+                    sensor_config["prim_path"] = f"/World/{sensor_config['name']}"
+                # Pop the desired position and orientation
+                local_position, local_orientation = sensor_config.pop("local_position", None), sensor_config.pop("local_orientation", None)
+                # Make sure sensor exists, grab its corresponding kwargs, and create the sensor
+                sensor = create_sensor(**sensor_config)
+                # Load an initialize this sensor
+                sensor.load()
+                sensor.initialize()
+                sensor.set_local_pose(local_position, local_orientation)
+                self._external_sensors[sensor.name] = sensor
+
+        assert og.sim.is_stopped(), "Simulator must be stopped after loading external sensors!"
+
     def _load_observation_space(self):
         # Grab robot(s) and task obs spaces
         obs_space = dict()
@@ -290,6 +305,15 @@ def _load_observation_space(self):
 
         # Also load the task obs space
         obs_space["task"] = self._task.load_observation_space()
+
+        # Also load any external sensors
+        if self._external_sensors is not None:
+            external_obs_space = dict()
+            for sensor_name, sensor in self._external_sensors.items():
+                # Load the sensor observation space
+                external_obs_space[sensor_name] = sensor.load_observation_space()
+            obs_space["external"] = gym.spaces.Dict(external_obs_space)
+
         return obs_space
 
     def load_observation_space(self):
@@ -339,6 +363,7 @@ def load(self):
         self._load_robots()
         self._load_objects()
         self._load_task()
+        self._load_external_sensors()
 
         og.sim.play()
         self.reset()
@@ -401,6 +426,13 @@ def get_obs(self):
         # Add task observations
         obs["task"] = self._task.get_obs(env=self)
 
+        # Add external sensor observations if they exist
+        if self._external_sensors is not None:
+            external_obs = dict()
+            for sensor_name, sensor in self._external_sensors.items():
+                external_obs[sensor_name] = sensor.get_obs()
+            obs["external"] = external_obs
+
         # Possibly flatten obs if requested
         if self._flatten_obs_space:
             obs = recursively_generate_flat_dict(dic=obs)
@@ -568,6 +600,15 @@ def robots(self):
         """
         return self.scene.robots
 
+    @property
+    def external_sensors(self):
+        """
+        Returns:
+            None or dict: If self.env_config["external_sensors"] is specified, returns the dict mapping sensor name to
+                instantiated sensor. Otherwise, returns None
+        """
+        return self._external_sensors
+
     @property
     def env_config(self):
         """
@@ -634,7 +675,14 @@ def default_config(self):
         return {
             # Environment kwargs
             "env": {
+                "action_timestep": 1 / 60.,
+                "physics_timestep": 1 / 60.,
+                "device": None,
+                "automatic_reset": False,
+                "flatten_action_space": False,
+                "flatten_obs_space": False,
                 "initial_pos_z_offset": 0.1,
+                "external_sensors": None,
             },
 
             # Rendering kwargs

omnigibson/examples/learning/navigation_policy_demo.py

@@ -119,20 +119,18 @@ def main():
     with open(f"{example_config_path}/turtlebot_nav.yaml", "r") as f:
         cfg = yaml.load(f, Loader=yaml.FullLoader)
 
+    # Make sure flattened obs and action space is used
+    cfg["env"]["flatten_action_space"] = True
+    cfg["env"]["flatten_obs_space"] = True
+
     # Only use RGB obs
     cfg["robots"][0]["obs_modalities"] = ["rgb"]
 
     # If we're not eval, turn off the start / goal markers so the agent doesn't see them
     if not args.eval:
         cfg["task"]["visualize_goal"] = False
 
-    env = og.Environment(
-        configs=cfg,
-        action_timestep=1 / 60.,
-        physics_timestep=1 / 60.,
-        flatten_action_space=True,
-        flatten_obs_space=True,
-    )
+    env = og.Environment(configs=cfg)
 
     # If we're evaluating, hide the ceilings and enable camera teleoperation so the user can easily
     # visualize the rollouts dynamically

omnigibson/examples/object_states/dicing_demo.py

@@ -71,7 +71,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     }
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1/60., physics_timestep=1/60.)
+    env = og.Environment(configs=cfg)
 
     # Grab reference to apple and knife
     apple = env.scene.object_registry("name", "apple")

omnigibson/examples/object_states/folded_unfolded_state_demo.py

@@ -57,7 +57,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     }
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1 / 60., physics_timestep=1 / 60.)
+    env = og.Environment(configs=cfg)
 
     # Grab object references
     carpet = env.scene.object_registry("name", "carpet")

omnigibson/examples/object_states/heat_source_or_sink_demo.py

@@ -30,7 +30,7 @@ def main():
     }
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1/60., physics_timestep=1/60.)
+    env = og.Environment(configs=cfg)
 
     # Get reference to stove object
     stove = env.scene.object_registry("name", "stove")

omnigibson/examples/object_states/heated_state_demo.py

@@ -40,7 +40,7 @@ def main():
     }
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1/60., physics_timestep=1/60.)
+    env = og.Environment(configs=cfg)
 
     # Set camera to appropriate viewing pose
     og.sim.viewer_camera.set_position_orientation(

omnigibson/examples/object_states/object_state_texture_demo.py

@@ -60,7 +60,7 @@ def main():
     }
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1/60., physics_timestep=1/60.)
+    env = og.Environment(configs=cfg)
 
     # Set camera to appropriate viewing pose
     og.sim.viewer_camera.set_position_orientation(

omnigibson/examples/object_states/onfire_demo.py

@@ -60,7 +60,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     }
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1 / 60., physics_timestep=1 / 60.)
+    env = og.Environment(configs=cfg)
 
     # Get reference to relevant objects
     stove = env.scene.object_registry("name", "stove")

omnigibson/examples/object_states/overlaid_demo.py

@@ -49,7 +49,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     }
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1 / 60., physics_timestep=1 / 60.)
+    env = og.Environment(configs=cfg)
 
     # Grab object references
     carpet = env.scene.object_registry("name", "carpet")

omnigibson/examples/object_states/particle_applier_remover_demo.py

@@ -132,7 +132,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     # object, so contacts will not be triggered with this object
 
     # Load the environment, then immediately stop the simulator since we need to add in the modifier object
-    env = og.Environment(configs=cfg, action_timestep=1/60., physics_timestep=1/60.)
+    env = og.Environment(configs=cfg)
     og.sim.stop()
 
     # Grab references to table

omnigibson/examples/object_states/particle_source_sink_demo.py

@@ -63,7 +63,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     cfg["objects"] = [sink_cfg]
 
     # Create the environment!
-    env = og.Environment(configs=cfg, action_timestep=1/60., physics_timestep=1/60.)
+    env = og.Environment(configs=cfg)
 
     # Set camera to ideal angle for viewing objects
     og.sim.viewer_camera.set_position_orientation(

omnigibson/examples/object_states/sample_kinematics_demo.py

@@ -93,7 +93,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     cfg["objects"] = objects_cfg
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1/60., physics_timestep=1/60.)
+    env = og.Environment(configs=cfg)
     env.step([])
 
     # Sample microwave and boxes

omnigibson/examples/object_states/slicing_demo.py

@@ -69,7 +69,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     }
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1/60., physics_timestep=1/60.)
+    env = og.Environment(configs=cfg)
 
     # Grab reference to apple and knife
     apple = env.scene.object_registry("name", "apple")

omnigibson/examples/object_states/temperature_demo.py

@@ -106,7 +106,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     }
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1 / 60., physics_timestep=1 / 60.)
+    env = og.Environment(configs=cfg)
 
     # Get reference to relevant objects
     stove = env.scene.object_registry("name", "stove")

omnigibson/examples/objects/draw_bounding_box.py

@@ -48,7 +48,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     }
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1/60., physics_timestep=1/60.)
+    env = og.Environment(configs=cfg)
 
     # Set camera to appropriate viewing pose
     cam = og.sim.viewer_camera

omnigibson/examples/objects/load_object_selector.py

@@ -54,7 +54,7 @@ def main(random_selection=False, headless=False, short_exec=False):
         cfg["scene"]["scene_model"] = "Rs_int"
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1 / 60., physics_timestep=1 / 60.)
+    env = og.Environment(configs=cfg)
 
     # Place the object so it rests on the floor
     obj = env.scene.object_registry("name", "obj")

omnigibson/examples/renderer_settings/renderer_settings_example.py

@@ -38,7 +38,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     }
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1/60., physics_timestep=1/60.)
+    env = og.Environment(configs=cfg)
 
     # Set camera to appropriate viewing pose
     cam = og.sim.viewer_camera