diff --git a/.gitignore b/.gitignore index 14ba4cc20d..26efb8d6f3 100644 --- a/.gitignore +++ b/.gitignore @@ -25,6 +25,8 @@ machines.txt *_vision_data tests/_fake_trajs tests/_fake_results +*.mp4 +video_frames/ # Jetbrains IDEs .idea/ diff --git a/predicators/behavior_utils/option_model_fns.py b/predicators/behavior_utils/option_model_fns.py index 6174af4862..82eedfbc02 100644 --- a/predicators/behavior_utils/option_model_fns.py +++ b/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])) diff --git a/predicators/envs/behavior.py b/predicators/envs/behavior.py index 7b4102d26e..82ef068d54 100644 --- a/predicators/envs/behavior.py +++ b/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()) diff --git a/predicators/main.py b/predicators/main.py index a828b9f3bf..de27151760 100644 --- a/predicators/main.py +++ b/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") diff --git a/predicators/settings.py b/predicators/settings.py index 919e7adf91..1c5118e6d3 100644 --- a/predicators/settings.py +++ b/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"]