Merge pull request #327 from StanfordVL/feat/control-updates

Feat/control updates

omnigibson/envs/env_base.py

@@ -195,6 +195,11 @@ def _load_scene(self):
         Load the scene and robot specified in the config file.
         """
         assert og.sim.is_stopped(), "Simulator must be stopped before loading scene!"
+
+        # Set the simulator settings
+        # NOTE: This must be done BEFORE the scene is loaded, or else all vision sensors can't retrieve observations
+        og.sim.set_simulation_dt(physics_dt=self.physics_timestep, rendering_dt=self.action_timestep)
+
         # Create the scene from our scene config
         scene = create_class_from_registry_and_config(
             cls_name=self.scene_config["type"],
@@ -204,8 +209,7 @@ def _load_scene(self):
         )
         og.sim.import_scene(scene)
 
-        # Set the simulator settings
-        og.sim.set_simulation_dt(physics_dt=self.physics_timestep, rendering_dt=self.action_timestep)
+        # Set the rendering settings
         og.sim.viewer_width = self.render_config["viewer_width"]
         og.sim.viewer_height = self.render_config["viewer_height"]
         og.sim.device = self.device

omnigibson/examples/robots/robot_control_example.py

@@ -90,7 +90,7 @@ def main(random_selection=False, headless=False, short_exec=False):
     cfg = dict(scene=scene_cfg, robots=[robot0_cfg])
 
     # Create the environment
-    env = og.Environment(configs=cfg, action_timestep=1/60., physics_timestep=1/60.)
+    env = og.Environment(configs=cfg, action_timestep=1/10., physics_timestep=1/60.)
 
     # Choose robot controller to use
     robot = env.robots[0]
@@ -131,9 +131,8 @@ def main(random_selection=False, headless=False, short_exec=False):
     step = 0
     while step != max_steps:
         action = action_generator.get_random_action() if control_mode == "random" else action_generator.get_teleop_action()
-        for _ in range(10):
-            env.step(action=action)
-            step += 1
+        env.step(action=action)
+        step += 1
 
     # Always shut down the environment cleanly at the end
     env.close()

omnigibson/objects/controllable_object.py

@@ -286,11 +286,12 @@ def _create_continuous_action_space(self):
 
     def apply_action(self, action):
         """
+        Converts inputted actions into low-level control signals
 
-        Converts inputted actions into low-level control signals and deploys them on the object
+        NOTE: This does NOT deploy control on the object. Use self.step() instead.
 
         Args:
-            n_array: n-DOF length array of actions to convert and deploy on the object
+            action (n-array): n-DOF length array of actions to apply to this object's internal controllers
         """
         # Store last action as the current action being applied
         self._last_action = action
@@ -304,25 +305,26 @@ def apply_action(self, action):
             self.action_dim, len(action)
         )
 
-        # Run convert actions to controls
-        control, control_type = self._actions_to_control(action=action)
-
-        # Deploy control signals
-        self.deploy_control(control=control, control_type=control_type, indices=None, normalized=False)
+        # First, loop over all controllers, and update the desired command
+        idx = 0
 
-    def _actions_to_control(self, action):
-        """
-        Converts inputted @action into low level control signals to deploy directly on the object.
-        This returns two arrays: the converted low level control signals and an array corresponding
-        to the specific ControlType for each signal.
+        for name, controller in self._controllers.items():
+            # Set command, then take a controller step
+            controller.update_command(command=action[idx : idx + controller.command_dim])
+            # Update idx
+            idx += controller.command_dim
 
-        Args:
-            action (n-array): n-DOF length array of actions to convert and deploy on the object
+        # If we haven't already created a physics callback, do so now so control gets updated every sim step
+        callback_name = f"{self.name}_controller_callback"
+        if not og.sim.physics_callback_exists(callback_name=callback_name):
+            og.sim.add_physics_callback(
+                callback_name=callback_name,
+                callback_fn=lambda x: self.step(),
+            )
 
-        Returns:
-            2-tuple:
-                - n-array: raw control signals to send to the object's joints
-                - list: control types for each joint
+    def step(self):
+        """
+        Takes a controller step across all controllers and deploys the computed control signals onto the object.
         """
         # First, loop over all controllers, and calculate the computed control
         control = dict()
@@ -332,8 +334,6 @@ def _actions_to_control(self, action):
         control_dict = self.get_control_dict()
 
         for name, controller in self._controllers.items():
-            # Set command, then take a controller step
-            controller.update_command(command=action[idx : idx + controller.command_dim])
             control[name] = {
                 "value": controller.step(control_dict=control_dict),
                 "type": controller.control_type,
@@ -350,8 +350,8 @@ def _actions_to_control(self, action):
             u_vec[idx] = ctrl["value"]
             u_type_vec[idx] = ctrl["type"]
 
-        # Return control
-        return u_vec, u_type_vec
+        # Deploy control signals
+        self.deploy_control(control=u_vec, control_type=u_type_vec, indices=None, normalized=False)
 
     def deploy_control(self, control, control_type, indices=None, normalized=False):
         """