Robust PLR

syllabus/curricula/plr/central_plr_wrapper.py

@@ -133,7 +133,7 @@ def __init__(
         self._gae_lambda = gae_lambda
         self._supress_usage_warnings = suppress_usage_warnings
         self._task2index = {task: i for i, task in enumerate(self.tasks)}
-        self._task_sampler = TaskSampler(self.tasks, action_space=action_space, **task_sampler_kwargs_dict)
+        self._task_sampler = TaskSampler(self.tasks, task_space=task_space, action_space=action_space, **task_sampler_kwargs_dict)
         self._rollouts = RolloutStorage(
             self._num_steps,
             self._num_processes,

syllabus/curricula/plr/plr_wrapper.py

@@ -149,6 +149,9 @@ class PrioritizedLevelReplay(Curriculum):
         gamma (float): The discount factor used to compute returns
         gae_lambda (float): The GAE lambda value.
         suppress_usage_warnings (bool): Whether to suppress warnings about improper usage.
+        robust_plr (bool): Option to use RobustPLR.
+        eval_envs: Evaluation environments for RobustPLR.
+        action_value_fn (callable): A function that takes an observation as input and returns an action and value.
         **curriculum_kwargs: Keyword arguments to pass to the curriculum.
     """
     REQUIRES_STEP_UPDATES = True
@@ -170,6 +173,9 @@ def __init__(
         suppress_usage_warnings=False,
         get_value=null,
         get_action_log_dist=null,
+        robust_plr: bool = False,  # Option to use RobustPLR
+        eval_envs = None,
+        action_value_fn = None,
         **curriculum_kwargs,
     ):
         # Preprocess curriculum intialization args
@@ -186,15 +192,22 @@ def __init__(
         task_sampler_kwargs_dict["num_actors"] = num_processes
         super().__init__(task_space, *curriculum_args, **curriculum_kwargs)
 
+        if robust_plr and eval_envs is None:
+            raise UsageError("RobustPLR requires evaluation environments to be provided.")
+
         self._num_steps = num_steps  # Number of steps stored in rollouts and used to update task sampler
         self._num_processes = num_processes  # Number of parallel environments
         self._gamma = gamma
         self._gae_lambda = gae_lambda
         self._supress_usage_warnings = suppress_usage_warnings
         self._get_action_log_dist = get_action_log_dist
         self._task2index = {task: i for i, task in enumerate(self.tasks)}
+        self._robust_plr = robust_plr
+        self._eval_envs = eval_envs
+        self.action_value_fn = action_value_fn
+
+        self._task_sampler = TaskSampler(self.tasks, task_space=task_space, action_space=action_space, robust_plr=robust_plr, eval_envs=eval_envs, action_value_fn = action_value_fn, **task_sampler_kwargs_dict)
 
-        self._task_sampler = TaskSampler(self.tasks, action_space=action_space, **task_sampler_kwargs_dict)
         self._rollouts = RolloutStorage(
             self._num_steps,
             self._num_processes,

syllabus/curricula/plr/storage.py

@@ -0,0 +1,69 @@
+import gymnasium as gym
+import torch
+
+class RolloutStorage(object):
+    def __init__(
+        self,
+        num_steps: int,
+        num_processes: int,
+        requires_value_buffers: bool,
+        action_space: gym.Space = None,
+    ):
+        self._requires_value_buffers = requires_value_buffers
+        self.tasks = torch.zeros(num_steps, num_processes, 1, dtype=torch.int)
+        self.masks = torch.ones(num_steps + 1, num_processes, 1)
+
+        if requires_value_buffers:
+            self.returns = torch.zeros(num_steps + 1, num_processes, 1)
+            self.rewards = torch.zeros(num_steps, num_processes, 1)
+            self.value_preds = torch.zeros(num_steps + 1, num_processes, 1)
+        else:
+            if action_space is None:
+                raise ValueError(
+                    "Action space must be provided to PLR for strategies 'policy_entropy', 'least_confidence', 'min_margin'"
+                )
+            self.action_log_dist = torch.zeros(num_steps, num_processes, action_space.n)
+
+        self.num_steps = num_steps
+        self.step = 0
+
+    def to(self, device):
+        self.masks = self.masks.to(device)
+        self.tasks = self.tasks.to(device)
+        if self._requires_value_buffers:
+            self.rewards = self.rewards.to(device)
+            self.value_preds = self.value_preds.to(device)
+            self.returns = self.returns.to(device)
+        else:
+            self.action_log_dist = self.action_log_dist.to(device)
+
+    def insert(self, masks, action_log_dist=None, value_preds=None, rewards=None, tasks=None):
+        if self._requires_value_buffers:
+            assert (value_preds is not None and rewards is not None), "Selected strategy requires value_preds and rewards"
+            if len(rewards.shape) == 3:
+                rewards = rewards.squeeze(2)
+            self.value_preds[self.step].copy_(torch.as_tensor(value_preds))
+            self.rewards[self.step].copy_(torch.as_tensor(rewards))
+            self.masks[self.step + 1].copy_(torch.as_tensor(masks))
+        else:
+            self.action_log_dist[self.step].copy_(action_log_dist)
+        if tasks is not None:
+            # assert isinstance(tasks[0], (int, torch.int32)), "Provided task must be an integer"
+            self.tasks[self.step].copy_(torch.as_tensor(tasks))
+        self.step = (self.step + 1) % self.num_steps
+
+    def after_update(self):
+        self.masks[0].copy_(self.masks[-1])
+
+    def compute_returns(self, next_value, gamma, gae_lambda):
+        assert self._requires_value_buffers, "Selected strategy does not use compute_rewards."
+        self.value_preds[-1] = next_value
+        gae = 0
+        for step in reversed(range(self.rewards.size(0))):
+            delta = (
+                self.rewards[step]
+                + gamma * self.value_preds[step + 1] * self.masks[step + 1]
+                - self.value_preds[step]
+            )
+            gae = delta + gamma * gae_lambda * self.masks[step + 1] * gae
+            self.returns[step] = gae + self.value_preds[step]