Save video (#78)

* started save-video

* update save-video

* save video works

* ran some checks

* fixed checks?

* ran checks again

* fix checks

* final touches

* revert changes in planning

* yapf

---------

Co-authored-by: Nishanth Kumar <[email protected]>

.gitignore

@@ -25,6 +25,8 @@ machines.txt
 *_vision_data
 tests/_fake_trajs
 tests/_fake_results
+*.mp4
+video_frames/
 
 # Jetbrains IDEs
 .idea/

predicators/behavior_utils/option_model_fns.py

@@ -69,21 +69,24 @@ def navigateToOptionModel(_init_state: State, env: "BehaviorEnv") -> None:
                          f"params {sample_arr}")
 
         if CFG.simulate_nav:
+            # Run a closed-loop controller for navigation.
+            curr_plan = plan[:]
             done_bit = False
             while not done_bit:
                 # Get expected position and orientation from plan.
-                expected_pos = np.array([plan[0][0], plan[0][1], robot_z])
+                expected_pos = np.array(
+                    [curr_plan[0][0], curr_plan[0][1], robot_z])
                 expected_orn = p.getQuaternionFromEuler(
-                    np.array([robot_orn[0], robot_orn[1], plan[0][2]]))
+                    np.array([robot_orn[0], robot_orn[1], curr_plan[0][2]]))
                 # In this case, we're at the final position we wanted to reach.
-                if len(plan) == 1:
+                if len(curr_plan) == 1:
                     done_bit = True
                     logging.info(
                         "PRIMITIVE: navigation policy completed execution!")
                 env.robots[0].set_position_orientation(expected_pos,
                                                        expected_orn)
                 env.step(np.zeros(env.action_space.shape))
-                plan.pop(0)
+                curr_plan.pop(0)
         target_pos = np.array([desired_xpos, desired_ypos, robot_z])
         target_orn = p.getQuaternionFromEuler(
             np.array([robot_orn[0], robot_orn[1], desired_zrot]))

predicators/envs/behavior.py

@@ -577,6 +577,9 @@ def set_igibson_behavior_env(self, task_num: int, task_instance_id: int,
         """Sets/resets the igibson_behavior_env."""
         np.random.seed(seed)
         env_creation_attempts = 0
+        save_video = CFG.env == "behavior" and CFG.behavior_save_video
+        if CFG.env == "behavior":
+            task_name = str(CFG.behavior_task_list)[2:-2]
         # NOTE: this while loop is necessary because in some cases
         # when CFG.randomize_init_state is True, creating a new
         # iGibson env may fail and we need to keep trying until
@@ -593,8 +596,10 @@ def set_igibson_behavior_env(self, task_num: int, task_instance_id: int,
                 instance_id=task_instance_id,
                 rng=self._rng,
             )
-            self.igibson_behavior_env.step(
-                np.zeros(self.igibson_behavior_env.action_space.shape))
+            self.igibson_behavior_env.step(np.zeros(
+                self.igibson_behavior_env.action_space.shape),
+                                           save_video=save_video,
+                                           task_name=task_name)
             ig_objs_bddl_scope = [
                 self._ig_object_name(obj) for obj in list(
                     self.igibson_behavior_env.task.object_scope.values())

predicators/main.py

@@ -52,6 +52,7 @@
 from predicators.approaches.gnn_approach import GNNApproach
 from predicators.datasets import create_dataset
 from predicators.envs import BaseEnv, create_new_env
+from predicators.envs.behavior import BehaviorEnv
 from predicators.planning import _run_plan_with_option_model
 from predicators.settings import CFG
 from predicators.structs import Dataset, InteractionRequest, \
@@ -122,6 +123,11 @@ def main() -> None:
     # Run the full pipeline.
     _run_pipeline(env, approach, stripped_train_tasks, offline_dataset)
     script_time = time.perf_counter() - script_start
+    if CFG.env == "behavior":  # pragma: no cover
+        assert isinstance(env, BehaviorEnv)
+        task_name = str(CFG.behavior_task_list)[2:-2]
+        env.igibson_behavior_env.simulator.viewer.make_video(
+            task_name=task_name)
     logging.info(f"\n\nMain script terminated in {script_time:.5f} seconds")
 
 

predicators/settings.py

@@ -120,6 +120,7 @@ class GlobalSettings:
         "wbm3_modifiable_full_obs.yaml",
     )
     behavior_mode = "headless"  # headless, pbgui, iggui
+    behavior_save_video = False  # True
     behavior_action_timestep = 1.0 / 10.0
     behavior_physics_timestep = 1.0 / 120.0
     behavior_task_list = ["re-shelving_library_books"]