Next task class for BatchedTask directly given by the TaskSampler

allenai · Jul 12, 2024 · 3363fee · 3363fee
1 parent f5afde3
commit 3363fee
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/allenact/base_abstractions/task.py b/allenact/base_abstractions/task.py
@@ -394,7 +394,6 @@ class BatchedTask(Generic[EnvType]):
 
     task_sampler: TaskSampler
     tasks: List[Task]
-    task_classes: List[type(Task)]
     callback_sensor_suite: Optional[SensorSuite]
 
     def __init__(
@@ -404,7 +403,7 @@ def __init__(
         task_info: Dict[str, Any],
         max_steps: int,
         task_sampler: TaskSampler,
-        task_classes: List[type(Task)],
+        task_class: type(Task),
         callback_sensor_suite: Optional[SensorSuite],
         **kwargs,
     ) -> None:
@@ -415,14 +414,13 @@ def __init__(
         # Keep a reference to the task sampler
         self.task_sampler = task_sampler
 
-        self.task_classes = task_classes
         self.callback_sensor_suite = callback_sensor_suite
 
         self.env = env
 
         # Instantiate the first actual task from the currently sampled info
         self.tasks = [
-            task_classes[0](
+            task_class(
                 env=env,
                 sensors=sensors,
                 task_info=task_info,

diff --git a/tests/make_it_batch/experiment.py b/tests/make_it_batch/experiment.py
@@ -197,7 +197,7 @@ def next_task(
                 max_steps=self.max_steps,
                 action_space=self._action_space,
                 task_sampler=self,
-                task_classes=[BatchableObjectNaviThorGridTask],
+                task_class=BatchableObjectNaviThorGridTask,
                 callback_sensor_suite=self.callback_sensor_suite,
             )
         else: