switch to mujoco env

allenact/base_abstractions/task.py

@@ -392,6 +392,7 @@ def __init__(
         max_steps: int,
         task_sampler: TaskSampler,
         task_classes: List[type(Task)],
+        state_views: List,
         callback_sensor_suite: Optional[SensorSuite],
         **kwargs,
     ) -> None:
@@ -407,6 +408,11 @@ def __init__(
         self.tasks = [task_classes[0](env=env, sensors=sensors, task_info=task_info, max_steps=max_steps, batch_index=0, **kwargs)]
         self.tasks[0].batch_index = 0
 
+        self.frames = None
+        self.depths = None
+        self.segs = None
+        self.state_views = state_views
+
         # If task_batch_size greater than 1, instantiate the rest of tasks (with task_batch_size set to 1)
         if self.task_sampler.task_batch_size > 1:
             for it in range(1, self.task_sampler.task_batch_size):
@@ -428,10 +434,22 @@ def observation_space(self):
 
     def get_observations(self, **kwargs) -> List[Any]:  #-> Dict[str, Any]:
         # Render all tasks in batch
-        self.tasks[0].env.render()  # assume this is stored locally in the env class
+        self.frames = self.render("rgb")
+        self.depths = self.render("depth")
+        self.segs = self.render("seg")
 
         # return {"batch_observations": [task.get_observations() for task in self.tasks]}
-        return [task.get_observations() for task in self.tasks]
+        return [task.get_observations(
+            frame=self.frames[idx],
+            depth=self.depths[idx],
+            seg=self.segs[idx],
+        ) for idx, task in enumerate(self.tasks)]
+
+    def update_state_views(self):
+        for idx, state_view in enumerate(self.state_views):
+            updated_state_view = self.tasks[0].env.call(state_view)
+            for idy, task in enumerate(self.tasks):
+                task.state_views[idx] = updated_state_view[idy]
 
     @property
     @abc.abstractmethod
@@ -456,7 +474,7 @@ def render(self, mode: str = "rgb", *args, **kwargs) -> np.ndarray:
 
         An numpy array corresponding to the requested render.
         """
-        raise NotImplementedError()
+        return self.tasks[0].env.render(mode=mode, *args, **kwargs)
 
     def step(self, action: Any) -> RLStepResult:
         srs = self._step(action=action)

tests/make_it_batch/batch_ai2thor_controller.py

@@ -1,32 +1,38 @@
 from typing import List
 from ai2thor.controller import Controller
+from allenact_plugins.ithor_plugin.ithor_environment import IThorEnvironment
 
 
 class BatchController:
     def __init__(
-        self,
-        task_batch_size: int,
-        **kwargs,
+            self,
+            task_batch_size: int,
+            **kwargs,
     ):
         self.task_batch_size = task_batch_size
-        self.controllers = [Controller(**kwargs) for _ in range(task_batch_size)]
+        self.controllers = [IThorEnvironment(**kwargs) for _ in range(task_batch_size)]
+        self._frames = []
 
     def step(self, actions: List[str]):
         assert len(actions) == self.task_batch_size
         for controller, action in zip(self.controllers, actions):
-            controller.step(action)
+            controller.step(action=action if action != "End" else "Pass")
+        self._frames = []
         return self.batch_last_event()
 
+    def get_agent_location(self):
+        return None
+
     def reset(
-        self,
-        idx: int,
-        scene: str,
+            self,
+            idx: int,
+            scene: str,
     ):
         self.controllers[idx].reset(scene)
 
     def batch_reset(
-        self,
-        scenes: List[str],
+            self,
+            scenes: List[str],
     ):
         for controller, scene in zip(self.controllers, scenes):
             controller.reset(scene)
@@ -42,7 +48,6 @@ def batch_last_event(self):
         return [controller.last_event for controller in self.controllers]
 
     def render(self):
-        frames = []
+        assert len(self._frames) == 0
         for controller in self.controllers:
-            frames.append(controller.last_event.frame)
-        return frames
+            self._frames.append(controller.last_event.frame)

tests/make_it_batch/experiment.py

@@ -31,55 +31,8 @@
 from allenact_plugins.ithor_plugin.ithor_environment import IThorEnvironment
 from allenact.base_abstractions.misc import RLStepResult
 
+from tests.make_it_batch.batch_ai2thor_controller import BatchController
 
-class BatchController:
-    def __init__(
-        self,
-        task_batch_size: int,
-        **kwargs,
-    ):
-        self.task_batch_size = task_batch_size
-        self.controllers = [IThorEnvironment(**kwargs) for _ in range(task_batch_size)]
-        self._frames = []
-
-    def step(self, actions: List[str]):
-        assert len(actions) == self.task_batch_size
-        for controller, action in zip(self.controllers, actions):
-            controller.step(action=action if action != "End" else "Pass")
-        self._frames = []
-        return self.batch_last_event()
-
-    def get_agent_location(self):
-        return None
-
-    def reset(
-        self,
-        idx: int,
-        scene: str,
-    ):
-        self.controllers[idx].reset(scene)
-
-    def batch_reset(
-        self,
-        scenes: List[str],
-    ):
-        for controller, scene in zip(self.controllers, scenes):
-            controller.reset(scene)
-
-    def stop(self):
-        for controller in self.controllers:
-            controller.stop()
-
-    def last_event(self, idx: int):
-        return self.controllers[idx].last_event
-
-    def batch_last_event(self):
-        return [controller.last_event for controller in self.controllers]
-
-    def render(self):
-        assert len(self._frames) == 0
-        for controller in self.controllers:
-            self._frames.append(controller.last_event.frame)
 
 
 class BatchableObjectNaviThorGridTask(ObjectNaviThorGridTask):