Adding default agentmode for ithor objectnav task

allenact_plugins/ithor_plugin/ithor_environment.py

@@ -38,11 +38,15 @@ def __init__(
         fov: float = FOV,
         player_screen_width: int = 300,
         player_screen_height: int = 300,
+        grid_size: float = 0.25,
+        rotate_step_degrees: int = 90,
         quality: str = "Very Low",
         restrict_to_initially_reachable_points: bool = False,
         make_agents_visible: bool = True,
         object_open_speed: float = 1.0,
         simplify_physics: bool = False,
+        snap_to_grid: bool = True,
+        **kwargs,
     ) -> None:
         """Initializer.
 
@@ -81,11 +85,13 @@ def __init__(
         self.controller: Optional[Controller] = None
         self._started = False
         self._quality = quality
+        self._snap_to_grid = snap_to_grid
 
         self._initially_reachable_points: Optional[List[Dict]] = None
         self._initially_reachable_points_set: Optional[Set[Tuple[float, float]]] = None
         self._move_mag: Optional[float] = None
-        self._grid_size: Optional[float] = None
+        self._grid_size: Optional[float] = grid_size
+        self._rotate_step_degrees = rotate_step_degrees
         self._visibility_distance = visibility_distance
         self._fov = fov
         self.restrict_to_initially_reachable_points = (
@@ -99,6 +105,9 @@ def __init__(
         self.start(None)
         # noinspection PyTypeHints
         self.controller.docker_enabled = docker_enabled  # type: ignore
+        self._extra_teleport_kwargs: Dict[
+            str, Any
+        ] = {}  # Used for backwards compatability with the teleport action
 
     @property
     def scene_name(self) -> str:
@@ -189,8 +198,12 @@ def start(
             width=self._start_player_screen_width,
             height=self._start_player_screen_height,
             local_executable_path=self._local_thor_build,
+            snapToGrid=self._snap_to_grid,
             quality=self._quality,
             server_class=ai2thor.fifo_server.FifoServer,
+            gridSize=self._grid_size,
+            rotateStepDegrees=self._rotate_step_degrees,
+            visibilityDistance=self._visibility_distance,
         )
 
         if (
@@ -728,6 +741,36 @@ def step(
 
         return sr
 
+    def set_object_filter(self, object_ids: List[str]):
+        self.controller.step("SetObjectFilter", objectIds=object_ids, renderImage=False)
+
+    def reset_object_filter(self):
+        self.controller.step("ResetObjectFilter", renderImage=False)
+
+    def teleport(
+        self,
+        pose: Dict[str, float],
+        rotation: Dict[str, float],
+        horizon: float = 0.0,
+    ):
+        try:
+            e = self.controller.step(
+                action="TeleportFull",
+                x=pose["x"],
+                y=pose["y"],
+                z=pose["z"],
+                rotation=rotation,
+                horizon=horizon,
+                **self._extra_teleport_kwargs,
+            )
+        except ValueError as e:
+            if len(self._extra_teleport_kwargs) == 0:
+                self._extra_teleport_kwargs["standing"] = True
+            else:
+                raise e
+            return self.teleport(pose=pose, rotation=rotation, horizon=horizon)
+        return e.metadata["lastActionSuccess"]
+
     @staticmethod
     def position_dist(
         p0: Mapping[str, Any],

allenact_plugins/ithor_plugin/ithor_task_samplers.py

@@ -1,6 +1,8 @@
 import copy
 import random
-from typing import List, Dict, Optional, Any, Union, cast
+import gzip
+import json
+from typing import List, Optional, Union, Dict, Any, cast, Tuple
 
 import gym
 
@@ -198,3 +200,218 @@ def set_seed(self, seed: int):
         self.seed = seed
         if seed is not None:
             set_seed(seed)
+
+
+class ObjectNavDatasetTaskSampler(TaskSampler):
+    def __init__(
+        self,
+        scenes: List[str],
+        scene_directory: str,
+        sensors: List[Sensor],
+        max_steps: int,
+        env_args: Dict[str, Any],
+        action_space: gym.Space,
+        # rewards_config: Dict,
+        seed: Optional[int] = None,
+        deterministic_cudnn: bool = False,
+        loop_dataset: bool = True,
+        allow_flipping=False,
+        env_class=IThorEnvironment,
+        **kwargs,
+    ) -> None:
+        # self.rewards_config = rewards_config
+        self.env_args = env_args
+        self.scenes = scenes
+        self.episodes = {
+            scene: ObjectNavDatasetTaskSampler.load_dataset(scene, scene_directory)
+            for scene in scenes
+        }
+        self.env_class = env_class
+        self.object_types = [
+            ep["object_type"] for scene in self.episodes for ep in self.episodes[scene]
+        ]
+        self.env: Optional[IThorEnvironment] = None
+        self.sensors = sensors
+        self.max_steps = max_steps
+        self._action_space = action_space
+        self.allow_flipping = allow_flipping
+        self.scene_counter: Optional[int] = None
+        self.scene_order: Optional[List[str]] = None
+        self.scene_id: Optional[int] = None
+        # get the total number of tasks assigned to this process
+        if loop_dataset:
+            self.max_tasks = None
+        else:
+            self.max_tasks = sum(len(self.episodes[scene]) for scene in self.episodes)
+        self.reset_tasks = self.max_tasks
+        self.scene_index = 0
+        self.episode_index = 0
+
+        self._last_sampled_task: Optional[ObjectNaviThorGridTask] = None
+
+        self.seed: Optional[int] = None
+        self.set_seed(seed)
+
+        if deterministic_cudnn:
+            set_deterministic_cudnn()
+
+        self.reset()
+
+    def _create_environment(self) -> IThorEnvironment:
+        env = self.env_class(
+            make_agents_visible=False,
+            object_open_speed=0.05,
+            restrict_to_initially_reachable_points=False,
+            **self.env_args,
+        )
+        return env
+
+    @staticmethod
+    def load_dataset(scene: str, base_directory: str) -> List[Dict]:
+        filename = (
+            "/".join([base_directory, scene])
+            if base_directory[-1] != "/"
+            else "".join([base_directory, scene])
+        )
+        filename += ".json.gz"
+        fin = gzip.GzipFile(filename, "r")
+        json_bytes = fin.read()
+        fin.close()
+        json_str = json_bytes.decode("utf-8")
+        data = json.loads(json_str)
+        random.shuffle(data)
+        return data
+
+    @staticmethod
+    def load_distance_cache_from_file(scene: str, base_directory: str) -> Dict:
+        filename = (
+            "/".join([base_directory, scene])
+            if base_directory[-1] != "/"
+            else "".join([base_directory, scene])
+        )
+        filename += ".json.gz"
+        fin = gzip.GzipFile(filename, "r")
+        json_bytes = fin.read()
+        fin.close()
+        json_str = json_bytes.decode("utf-8")
+        data = json.loads(json_str)
+        return data
+
+    @property
+    def __len__(self) -> Union[int, float]:
+        """Length.
+
+        # Returns
+
+        Number of total tasks remaining that can be sampled. Can be float('inf').
+        """
+        return float("inf") if self.max_tasks is None else self.max_tasks
+
+    @property
+    def total_unique(self) -> Optional[Union[int, float]]:
+        return self.reset_tasks
+
+    @property
+    def last_sampled_task(self) -> Optional[ObjectNaviThorGridTask]:
+        return self._last_sampled_task
+
+    def close(self) -> None:
+        if self.env is not None:
+            self.env.stop()
+
+    @property
+    def all_observation_spaces_equal(self) -> bool:
+        """Check if observation spaces equal.
+
+        # Returns
+
+        True if all Tasks that can be sampled by this sampler have the
+            same observation space. Otherwise False.
+        """
+        return True
+
+    @property
+    def length(self) -> Union[int, float]:
+        """Length.
+
+        # Returns
+
+        Number of total tasks remaining that can be sampled. Can be float('inf').
+        """
+        return float("inf") if self.max_tasks is None else self.max_tasks
+
+    def next_task(
+        self, force_advance_scene: bool = False
+    ) -> Optional[ObjectNaviThorGridTask]:
+        if self.max_tasks is not None and self.max_tasks <= 0:
+            return None
+
+        if self.episode_index >= len(self.episodes[self.scenes[self.scene_index]]):
+            self.scene_index = (self.scene_index + 1) % len(self.scenes)
+            # shuffle the new list of episodes to train on
+            random.shuffle(self.episodes[self.scenes[self.scene_index]])
+            self.episode_index = 0
+        scene = self.scenes[self.scene_index]
+        episode = self.episodes[scene][self.episode_index]
+        if self.env is None:
+            self.env = self._create_environment()
+
+        if scene.replace("_physics", "") != self.env.scene_name.replace("_physics", ""):
+            self.env.reset(scene_name=scene)
+        else:
+            self.env.reset_object_filter()
+
+        self.env.set_object_filter(
+            object_ids=[
+                o["objectId"]
+                for o in self.env.last_event.metadata["objects"]
+                if o["objectType"] == episode["object_type"]
+            ]
+        )
+
+        task_info = {"scene": scene, "object_type": episode["object_type"]}
+        if len(task_info) == 0:
+            get_logger().warning(
+                "Scene {} does not contain any"
+                " objects of any of the types {}.".format(scene, self.object_types)
+            )
+        task_info["initial_position"] = episode["initial_position"]
+        task_info["initial_orientation"] = episode["initial_orientation"]
+        task_info["initial_horizon"] = episode.get("initial_horizon", 0)
+        task_info["distance_to_target"] = episode.get("shortest_path_length")
+        task_info["path_to_target"] = episode.get("shortest_path")
+        task_info["object_type"] = episode["object_type"]
+        task_info["id"] = episode["id"]
+        if self.allow_flipping and random.random() > 0.5:
+            task_info["mirrored"] = True
+        else:
+            task_info["mirrored"] = False
+
+        self.episode_index += 1
+        if self.max_tasks is not None:
+            self.max_tasks -= 1
+        if not self.env.teleport(
+            pose=episode["initial_position"],
+            rotation=episode["initial_orientation"],
+            horizon=episode.get("initial_horizon", 0),
+        ):
+            return self.next_task()
+        self._last_sampled_task = ObjectNaviThorGridTask(
+            env=self.env,
+            sensors=self.sensors,
+            task_info=task_info,
+            max_steps=self.max_steps,
+            action_space=self._action_space,
+        )
+
+        return self._last_sampled_task
+
+    def reset(self):
+        self.episode_index = 0
+        self.scene_index = 0
+        self.max_tasks = self.reset_tasks
+
+    def set_seed(self, seed: int):
+        self.seed = seed
+        if seed is not None:
+            set_seed(seed)