comrad-bench 1.0.0__py3-none-any.whl

comrad/curriculum/__init__.py

@@ -0,0 +1,30 @@
+from comrad.curriculum.base import Curriculum
+from comrad.curriculum.uniform import UniformCurriculum
+from comrad.curriculum.sequential import SequentialCurriculum
+from comrad.curriculum.learning_progress import LearningProgress
+from comrad.curriculum.omni import OMNICurriculum
+from comrad.curriculum.plr import PrioritizedLevelReplay
+
+def make_curriculum(n_tasks: int, strategy: str = "uniform", **kwargs) -> Curriculum:
+    """Factory to spin up proper curriculum based on configuration."""
+    strategies = {
+        "uniform": UniformCurriculum,
+        "sequential": SequentialCurriculum,
+        "learning_progress": LearningProgress,
+        "omni": OMNICurriculum,
+        "plr": PrioritizedLevelReplay,
+    }
+    
+    assert strategy in strategies, f"Unknown strategy '{strategy}'. Available: {list(strategies.keys())}"
+    
+    return strategies[strategy](n_tasks=n_tasks, **kwargs)
+
+__all__ = [
+    "Curriculum",
+    "UniformCurriculum",
+    "SequentialCurriculum",
+    "LearningProgress",
+    "OMNICurriculum",
+    "PrioritizedLevelReplay",
+    "make_curriculum"
+]

comrad/curriculum/base.py

@@ -0,0 +1,133 @@
+import multiprocessing
+import os
+from typing import Optional
+import numpy as np
+
+
+class Curriculum:
+    """
+    Abstract base curriculum class for a fixed discrete batch of tasks (WADs).
+    Requires multiprocessing arrays and locks for disjoint worker/learner access.
+    """
+
+    EMA_ALPHA = 0.1
+    SCORE_MEAN_VALUE_L1   = "mean_value_l1"
+    SCORE_MAX_VALUE_L1    = "max_value_l1"
+    SCORE_MEAN_ADVANTAGE  = "mean_advantage"
+    SCORE_MAX_ADVANTAGE   = "max_advantage"
+    SCORE_MEAN_ENTROPY    = "mean_entropy"
+    SCORE_MAX_ENTROPY     = "max_entropy"
+
+    def __init__(self, n_tasks: int, seed: Optional[int] = None, tasks: Optional[list] = None, **kwargs):
+        self.ctx = multiprocessing.get_context("spawn")
+        self._n = n_tasks
+        self._base_seed = int(seed) if seed is not None else int.from_bytes(os.urandom(8), "little")
+        self._rng_counter = self.ctx.Value("Q", 0)
+        self._task_ids = tuple(tasks) if tasks is not None else None
+
+        self._weights = self.ctx.Array('d', [1.0 / n_tasks] * n_tasks)
+        self._counts = self.ctx.Array('l', [0] * n_tasks)
+        self._lock = self.ctx.Lock()
+
+    def update(self, task_idx: int, episode_return: float) -> None:
+        """Called at end of each episode."""
+        with self._lock:
+            self._counts[task_idx] += 1
+            self._update_logic(task_idx, float(episode_return))
+            self._recompute_weights()
+
+    def update_task_score(self, task_idx: int, score: dict, num_steps: int = 1) -> None:
+        """Used by sample factory learners (e.g. PLR)"""
+        pass
+
+    def sample(self) -> int:
+        with self._lock:
+            task_idx = self._sample_logic()
+        return task_idx
+
+    def _sample_logic(self) -> int:
+        """Isolated locked core for actually rolling the random sample"""
+        weights = np.array(self._weights[:])
+        return self._choice(self._n, p=weights)
+
+    def _next_rng(self) -> np.random.Generator:
+        """
+        Generate randomness from a single shared draw counter so spawned workers
+        do not each replay an identical local RNG stream.
+        """
+        counter = self._rng_counter.value
+        self._rng_counter.value = counter + 1
+        seed = (self._base_seed + 0x9E3779B97F4A7C15 * counter) & ((1 << 64) - 1)
+        return np.random.default_rng(seed)
+
+    def _choice(self, n: int, p=None) -> int:
+        return int(self._next_rng().choice(n, p=p))
+
+    def _random(self) -> float:
+        return float(self._next_rng().random())
+
+    def _update_logic(self, task_idx: int, episode_return: float) -> None:
+        """Sub-classes override this to implement update behavior."""
+        pass
+
+    def _recompute_weights(self) -> None:
+        """Sub-classes override this to update probabilities inside self._weights."""
+        pass
+
+    def task_weights(self) -> np.ndarray:
+        with self._lock:
+            return np.array(self._weights[:])
+
+    def task_returns(self) -> np.ndarray:
+        """Return fast moving average or appropriate returns proxy."""
+        # By default just returns zeroes, overridden in LP/OMNI classes
+        return np.zeros(self._n)
+
+    def task_scores(self) -> np.ndarray:
+        """PLR scores or similar metric."""
+        return np.zeros(self._n)
+
+    def metrics(self) -> dict:
+        """General metrics reporting."""
+        return {}
+
+    def state_dict(self) -> dict:
+        with self._lock:
+            state = {
+                "version": 1,
+                "curriculum_type": type(self).__name__,
+                "n_tasks": self._n,
+                "base_seed": int(self._base_seed),
+                "rng_counter": int(self._rng_counter.value),
+                "weights": list(self._weights[:]),
+                "counts": list(self._counts[:]),
+            }
+            if self._task_ids is not None:
+                state["task_ids"] = list(self._task_ids)
+            state.update(self._state_dict_locked())
+            return state
+
+    def load_state_dict(self, state: dict) -> None:
+        expected_type = type(self).__name__
+        actual_type = state.get("curriculum_type", expected_type)
+        if actual_type != expected_type:
+            raise ValueError(f"Expected curriculum state for {expected_type}, got {actual_type}")
+        if int(state.get("n_tasks", self._n)) != self._n:
+            raise ValueError(f"Expected curriculum state for {self._n} tasks, got {state.get('n_tasks')}")
+        expected_task_ids = list(self._task_ids) if self._task_ids is not None else None
+        actual_task_ids = state.get("task_ids", None)
+        if expected_task_ids is not None and actual_task_ids != expected_task_ids:
+            raise ValueError(f"Expected curriculum task ids {expected_task_ids}, got {actual_task_ids}")
+
+        with self._lock:
+            self._base_seed = int(state.get("base_seed", self._base_seed))
+            self._rng_counter.value = int(state.get("rng_counter", 0))
+            self._weights[:] = [float(weight) for weight in state.get("weights", self._weights[:])]
+            self._counts[:] = [int(count) for count in state.get("counts", self._counts[:])]
+            self._load_state_dict_locked(state)
+
+    def _state_dict_locked(self) -> dict:
+        return {}
+
+    def _load_state_dict_locked(self, state: dict) -> None:
+        pass

comrad/curriculum/learning_progress.py

@@ -0,0 +1,136 @@
+import numpy as np
+from comrad.curriculum.base import Curriculum
+
+class LearningProgress(Curriculum):
+    """
+    Syllabus-inspired Learning Progress curriculum tracking fast and slow EMAs over
+    normalized per-task success metrics. The raw terminal `true_objective` signal is
+    tracked separately for diagnostics/checkpoint continuity.
+    """
+
+    def __init__(
+        self,
+        n_tasks: int,
+        p_theta: float = 0.1,
+        max_return: float = 100.0,
+        min_return: float = -100.0,
+        uniform_prob: float = 0.25,
+        **kwargs,
+    ):
+        super().__init__(n_tasks, **kwargs)
+        self.p_theta = p_theta
+        self.max_return = max_return
+        self.min_return = min_return
+        self.uniform_prob = uniform_prob
+        self._stale_dist = self.ctx.Value('b', True)
+        self._p_fast = self.ctx.Array('d', [0.0] * n_tasks)
+        self._p_slow = self.ctx.Array('d', [0.0] * n_tasks)
+        self._p_true = self.ctx.Array('d', [0.0] * n_tasks)
+        self._p_seen = self.ctx.Array('d', [0.0] * n_tasks)  # 0.0 = unseen, 1.0 = seen
+
+    def _sample_logic(self) -> int:
+        if self._stale_dist.value:
+            self._recompute_weights()
+            self._stale_dist.value = False
+        weights = np.array(self._weights[:])
+        return self._choice(self._n, p=weights)
+
+    def update(self, task_idx: int, task_score: float) -> None:
+        with self._lock:
+            self._counts[task_idx] += 1
+            self._update_logic(task_idx, float(task_score))
+            self._stale_dist.value = True
+
+    def _normalize(self, task_score: float) -> float:
+        """Map raw `true_objective` success metrics to [0, 1]."""
+        r_min = self.min_return
+        r_max = self.max_return
+        if r_max <= r_min:
+            return 0.0
+        normalized = (task_score - r_min) / (r_max - r_min)
+        return float(np.clip(normalized, 0.0, 1.0))
+
+    def _update_logic(self, task_idx: int, task_score: float) -> None:
+        """Update normalized LP EMAs plus a raw `true_objective` EMA."""
+        p = self._normalize(task_score)
+        alpha = self.EMA_ALPHA
+
+        if self._p_seen[task_idx] == 0.0:
+            self._p_fast[task_idx] = p
+            self._p_slow[task_idx] = p
+            self._p_true[task_idx] = task_score
+            self._p_seen[task_idx] = 1.0
+        else:
+            old_fast = self._p_fast[task_idx]
+            old_true = self._p_true[task_idx]
+            self._p_fast[task_idx] = p * alpha + old_fast * (1.0 - alpha)
+            self._p_slow[task_idx] = self._p_fast[task_idx] * alpha + self._p_slow[task_idx] * (1.0 - alpha)
+            self._p_true[task_idx] = task_score * alpha + old_true * (1.0 - alpha)
+
+    def _reweight(self, p: np.ndarray) -> np.ndarray:
+        numerator = p * (1.0 - self.p_theta)
+        denominator = p + self.p_theta * (1.0 - 2.0 * p)
+        return numerator / denominator
+
+    def _recompute_weights(self) -> None:
+        seen = np.array(self._p_seen[:]) > 0.0
+        unseen = ~seen
+
+        # LP is undefined until every task has at least some signal.
+        # Keep the curriculum uniform during this cold-start phase instead of
+        # forcing unseen-only sweeps or collapsing onto the first updated task.
+        if unseen.any():
+            self._weights[:] = (np.ones(self._n) / self._n).tolist()
+            return
+
+        fast = np.array(self._p_fast[:])
+        slow = np.array(self._p_slow[:])
+        # Compute learning progress as absolute difference of reweighted EMAs
+        learning_progress = np.abs(self._reweight(fast) - self._reweight(slow))
+        # Keep tasks with any normalized success signal eligible when LP is flat.
+        # `_p_true` stores raw objectives for OMNI ordering, so using its sign here
+        # would incorrectly bias sampling toward tasks whose reward scale is positive.
+        posidxs = [i for i, lp in enumerate(learning_progress) if lp > 0.0 or fast[i] > 0.0]
+        any_progress = len(posidxs) > 0
+        subprobs = learning_progress[posidxs] if any_progress else learning_progress
+
+        std = np.std(subprobs)
+        # z-score
+        subprobs = (subprobs - np.mean(subprobs)) / (std if std else 1.0)
+        # sigmoid
+        subprobs = 1.0 / (1.0 + np.exp(-subprobs))
+        # normalize
+        subprobs = subprobs / (np.sum(subprobs))
+
+        if any_progress:
+            task_dist = np.zeros(self._n)
+            task_dist[posidxs] = subprobs
+        else:
+            task_dist = subprobs
+
+        if self.uniform_prob > 0.0:
+            uniform = np.ones(self._n) / self._n
+            task_dist = (1.0 - self.uniform_prob) * task_dist + self.uniform_prob * uniform
+
+        task_dist = task_dist / np.sum(task_dist)
+        self._weights[:] = task_dist.tolist()
+
+    def task_returns(self) -> np.ndarray:
+        with self._lock:
+            return np.array(self._p_fast[:])
+
+    def _state_dict_locked(self) -> dict:
+        return {
+            "stale_dist": bool(self._stale_dist.value),
+            "p_fast": list(self._p_fast[:]),
+            "p_slow": list(self._p_slow[:]),
+            "p_true": list(self._p_true[:]),
+            "p_seen": list(self._p_seen[:]),
+        }
+
+    def _load_state_dict_locked(self, state: dict) -> None:
+        self._stale_dist.value = bool(state.get("stale_dist", True))
+        self._p_fast[:] = [float(value) for value in state.get("p_fast", self._p_fast[:])]
+        self._p_slow[:] = [float(value) for value in state.get("p_slow", self._p_slow[:])]
+        self._p_true[:] = [float(value) for value in state.get("p_true", self._p_true[:])]
+        self._p_seen[:] = [float(value) for value in state.get("p_seen", self._p_seen[:])]

comrad/curriculum/observer.py

@@ -0,0 +1,100 @@
+import json
+import os
+from os.path import join
+
+import numpy as np
+import torch
+
+from sample_factory.algo.runners.runner import AlgoObserver
+from sample_factory.algo.learning.learner import Learner
+from sample_factory.utils.utils import experiment_dir, log
+
+
+def curriculum_state_file(cfg) -> str:
+    return join(experiment_dir(cfg), "curriculum_state.json")
+
+
+def save_curriculum_state(cfg, curriculum) -> None:
+    if curriculum is None:
+        return
+
+    state_path = curriculum_state_file(cfg)
+    tmp_path = f"{state_path}.tmp"
+    with open(tmp_path, "w") as state_file:
+        json.dump(curriculum.state_dict(), state_file, indent=2)
+    os.replace(tmp_path, state_path)
+
+
+def load_curriculum_state(cfg, curriculum) -> bool:
+    if curriculum is None:
+        return False
+
+    state_path = curriculum_state_file(cfg)
+    if not os.path.isfile(state_path):
+        return False
+
+    with open(state_path, "r") as state_file:
+        state = json.load(state_file)
+
+    curriculum.load_state_dict(state)
+    return True
+
+
+class CurriculumObserver(AlgoObserver):
+    def __init__(self, curriculum):
+        self.curriculum = curriculum
+
+    def on_init(self, runner) -> None:
+        return None
+
+    def on_training_step(self, runner, training_iteration_since_resume: int) -> None:
+        return None
+
+    def on_stop(self, runner) -> None:
+        return None
+
+    def extra_summaries(self, runner, policy_id, writer, env_steps) -> None:
+        if policy_id != 0 or self.curriculum is None:
+            return
+
+        weights = self.curriculum.task_weights()
+        if weights.size > 0:
+            nonzero_weights = weights[weights > 0]
+            if nonzero_weights.size > 0:
+                entropy = float(-(nonzero_weights * np.log(nonzero_weights)).sum())
+                writer.add_scalar("curriculum/entropy", entropy, env_steps)
+            writer.add_scalar("curriculum/nonzero_probability_tasks", int(np.count_nonzero(weights)), env_steps)
+
+        task_returns = self.curriculum.task_returns()
+        if task_returns.size > 0 and not np.allclose(task_returns, 0.0):
+            writer.add_scalar("curriculum/mean_task_return", float(np.mean(task_returns)), env_steps)
+
+        metrics = self.curriculum.metrics()
+        for key, value in metrics.items():
+            if np.isscalar(value):
+                writer.add_scalar(f"curriculum/{key}", float(value), env_steps)
+            elif isinstance(value, np.ndarray) and value.size > 0 and np.issubdtype(value.dtype, np.number):
+                writer.add_scalar(f"curriculum/mean_{key}", float(np.mean(value)), env_steps)
+
+
+def restore_curriculum_state(cfg, curriculum) -> None:
+    if curriculum is None or cfg.restart_behavior != "resume":
+        return
+
+    checkpoint_dir = Learner.checkpoint_dir(cfg, 0)
+    load_kind = getattr(cfg, "load_checkpoint_kind", "latest")
+    pattern = f"{dict(latest='checkpoint', best='best')[load_kind]}_*"
+    checkpoints = Learner.get_checkpoints(checkpoint_dir, pattern=pattern)
+    if not checkpoints:
+        return
+
+    checkpoint_dict = Learner.load_checkpoint(checkpoints, torch.device("cpu"))
+    if checkpoint_dict is None:
+        return
+
+    state = checkpoint_dict.get("curriculum_state", None)
+    if state is None:
+        raise RuntimeError("Checkpoint is missing curriculum_state; cannot safely resume curriculum-enabled training")
+
+    curriculum.load_state_dict(state)
+    log.info("Restored curriculum state from checkpoint %s", checkpoints[-1])

comrad/curriculum/omni.py

@@ -0,0 +1,74 @@
+import numpy as np
+from typing import Optional
+from comrad.curriculum.learning_progress import LearningProgress
+
+class OMNICurriculum(LearningProgress):
+    """
+    Syllabus-inspired OMNI curriculum tracking fast and slow EMAs masked by tasks dependencies
+    suppressing mathematically boring tasks once prerequisites are met.
+    """
+
+    def __init__(self, n_tasks: int, tasks: list, interestingness: Optional[dict] = None, **kwargs):
+        super().__init__(n_tasks, tasks=tasks, **kwargs)
+        self.tasks = tasks or list(range(n_tasks))
+        if interestingness is not None:
+            bad = set(interestingness.keys()) - set(self.tasks)
+            if bad:
+                raise ValueError(
+                    f"interestingness keys not in tasks: {bad}. "
+                    f"Use task identifiers only."
+            )
+            self._interestingness = {
+                task: {other: bool(is_interesting) for other, is_interesting in mapping.items()}
+                for task, mapping in interestingness.items()
+            }
+        else:
+            # Default: mastery of a task makes itself boring, but other tasks remain interesting.
+            self._interestingness = {
+                t: {other: (other != t) for other in self.tasks}
+                for t in self.tasks
+            }
+
+    def _recompute_weights(self) -> None:
+        """Combine LP weight calculation with "boring" penalty mask."""
+        super()._recompute_weights() # Update weights purely by LP
+
+        lp_dist = np.array(self._weights[:])
+        mastery = np.array(self._p_fast[:])
+
+        # OMNI prerequisite masking must be ordered by a comparable mastery signal.
+        # Raw true_objective values can differ wildly across tasks, so use the
+        # normalized fast EMA from LearningProgress instead.
+        tasks_by_success = np.argsort(mastery)
+        interesting_tasks = set()
+        boring_tasks = set()
+
+        for task_idx in tasks_by_success[::-1]:
+            if task_idx in interesting_tasks or task_idx in boring_tasks:
+                continue
+            interesting_tasks.add(task_idx)
+
+            task = self.tasks[task_idx]
+            if task not in self._interestingness:
+                continue
+
+            for other_idx, other_task in enumerate(self.tasks):
+                if other_task not in self._interestingness[task]:
+                    boring_tasks.add(other_idx)
+                    continue
+                if (
+                    other_idx not in interesting_tasks
+                    and other_idx not in boring_tasks
+                    and not self._interestingness[task][other_task]
+                ):
+                    boring_tasks.add(other_idx)
+
+        # Build weight mask where boring items are heavily penalized
+        moi_weight = np.ones(self._n)
+        for i in boring_tasks:
+            moi_weight[i] = 0.001
+
+        dist = lp_dist * moi_weight
+        norm_dist = dist / dist.sum() if dist.sum() > 0 else np.ones(self._n) / self._n
+
+        self._weights[:] = norm_dist.tolist()

comrad/curriculum/plr.py

@@ -0,0 +1,208 @@
+import numpy as np
+from comrad.curriculum.base import Curriculum
+
+class PrioritizedLevelReplay(Curriculum):
+    """
+    Syllabus-inspired Prioritized Level Replay (PLR) tracking PPO rollout metrics natively.
+    Replays tasks proportional to learning potential metrics and random unseen tasks.
+    """
+
+    def __init__(
+        self,
+        n_tasks: int,
+        staleness_coef: float = 0.1,
+        score_transform: str = "rank",
+        temperature: float = 0.1,
+        staleness_transform: str = "power",
+        staleness_temperature: float = 1.0,
+        plr_score_key: str = "mean_value_l1",
+        alpha: float = 1.0,
+        replay_schedule: str = "proportionate",
+        replay_prob: float = 0.5,
+        rho: float = 1.0,
+        max_score_coef: float = 0.0,
+        eps: float = 0.05,
+        **kwargs
+    ):
+        super().__init__(n_tasks, **kwargs)
+        self._staleness_coef = staleness_coef
+        self._score_transform = score_transform
+        self._temperature = temperature
+        self._staleness_transform = staleness_transform
+        self._staleness_temperature = staleness_temperature
+        self._plr_score_key = plr_score_key
+        self._alpha = alpha
+        self._replay_schedule = replay_schedule
+        self._replay_prob = replay_prob
+        self._rho = rho
+        self._max_score_coef = max_score_coef
+        self._eps = eps
+
+        self._task_scores = self.ctx.Array('d', [0.0] * n_tasks)
+        self._task_score_steps = self.ctx.Array('l', [0] * n_tasks)
+        self._task_staleness = self.ctx.Array('d', [0.0] * n_tasks)
+        self._unseen_task_weights = self.ctx.Array('d', [1.0] * n_tasks)
+
+    def _update_logic(self, task_idx: int, episode_return: float) -> None:
+        """Called conventionally by workers. PLR relies on central PPO learner metrics natively."""
+        pass
+
+    def update_task_score(self, task_idx: int, score: dict, num_steps: int = 1) -> None:
+        """Inject from sample factory using native PLR learner loops."""
+        raw = float(score.get(self._plr_score_key, 0.0))
+        raw_max = float(score.get(self._plr_score_key.replace("mean_", "max_"), raw))
+        total_score = self._max_score_coef * raw_max + (1.0 - self._max_score_coef) * raw
+        num_steps = max(int(num_steps), 1)
+
+        with self._lock:
+            old = self._task_scores[task_idx]
+            self._task_scores[task_idx] = (1.0 - self._alpha) * old + self._alpha * total_score
+            self._task_score_steps[task_idx] += num_steps
+            self._unseen_task_weights[task_idx] = 0.0
+            self._recompute_weights()
+
+    def _sample_logic(self) -> int:
+        """Override to implement complex replay schedule logic directly inside lock."""
+        proportion_seen = self._proportion_seen()
+
+        if proportion_seen < self._rho:
+            unseen = np.array(self._unseen_task_weights[:])
+            if (unseen > 0).any():
+                return self._sample_unseen_level(unseen)
+            return self._sample_replay_level() # fallback to replay if all seen but still under rho
+
+        do_replay = self._sample_replay_decision(proportion_seen)
+        if do_replay:
+            return self._sample_replay_level()
+        else:
+            unseen = np.array(self._unseen_task_weights[:])
+            if (unseen > 0).any():
+                return self._sample_unseen_level(unseen)
+            return self._sample_replay_level()
+
+    def _proportion_seen(self) -> float:
+        unseen = np.array(self._unseen_task_weights[:])
+        return (self._n - (unseen > 0).sum()) / self._n
+
+    def _sample_replay_decision(self, proportion_seen: float) -> bool:
+        if self._replay_schedule == "fixed":
+            if proportion_seen >= self._rho:
+                if self._random() < self._replay_prob or proportion_seen >= 1.0:
+                    return True
+            return False
+        else:  # proportionate
+            if proportion_seen >= self._rho and self._random() < min(proportion_seen, self._replay_prob):
+                return True
+            return False
+
+    def _sample_replay_level(self) -> int:
+        weights = self._calculate_score_weights()
+        task_idx = self._choice(self._n, p=weights)
+        self._update_staleness(task_idx)
+        return task_idx
+
+    def _sample_unseen_level(self, unseen: np.ndarray) -> int:
+        unseen_w = unseen / unseen.sum()
+        task_idx = self._choice(self._n, p=unseen_w)
+        self._update_staleness(task_idx)
+        return task_idx
+
+    def _update_staleness(self, selected_idx: int) -> None:
+        if self._staleness_coef > 0:
+            staleness = np.array(self._task_staleness[:])
+            staleness += 1.0
+            staleness[selected_idx] = 0.0
+            self._task_staleness[:] = staleness.tolist()
+            self._recompute_weights()
+
+    def _calculate_score_weights(self) -> np.ndarray:
+        scores = np.array(self._task_scores[:])
+        unseen = np.array(self._unseen_task_weights[:])
+        staleness = np.array(self._task_staleness[:])
+
+        weights = self._apply_score_transform(self._score_transform, self._temperature, scores)
+        weights *= (1.0 - unseen)
+        z = weights.sum()
+        if z > 0:
+            weights /= z
+        else:
+            seen = (unseen == 0.0)
+            weights = seen.astype(float)
+            if seen.any():
+                weights /= weights.sum()
+
+        if self._staleness_coef > 0:
+            staleness_w = self._apply_score_transform(
+                self._staleness_transform, self._staleness_temperature, staleness
+            )
+            staleness_w *= (1.0 - unseen)
+            z_s = staleness_w.sum()
+            if z_s > 0:
+                staleness_w /= z_s
+            else:
+                staleness_w = (1.0 / self._n) * (1.0 - unseen)
+            weights = (1.0 - self._staleness_coef) * weights + self._staleness_coef * staleness_w
+
+        total = weights.sum()
+        return weights / total if total > 0 else weights
+
+    def _recompute_weights(self) -> None:
+        self._weights[:] = self._calculate_score_weights().tolist()
+
+    def _apply_score_transform(self, transform: str, temperature: float, scores: np.ndarray) -> np.ndarray:
+        if transform == "rank":
+            temp = np.flip(scores.argsort())
+            ranks = np.empty_like(temp, dtype=float)
+            ranks[temp] = np.arange(len(temp)) + 1
+            return 1.0 / (ranks ** (1.0 / temperature))
+        elif transform == "power":
+            eps = 0.0 if self._staleness_coef > 0 else 1e-3
+            return (np.clip(scores, 0, None) + eps) ** (1.0 / temperature)
+        elif transform == "softmax":
+            return np.exp(scores / temperature)
+        elif transform == "max":
+            w = np.zeros_like(scores)
+            scores_ = scores.copy()
+            scores_[np.array(self._unseen_task_weights[:]) > 0] = -np.inf
+            argmax_candidates = np.flatnonzero(np.isclose(scores_, scores_.max()))
+            argmax = argmax_candidates[self._choice(len(argmax_candidates))]
+            w[argmax] = 1.0
+            return w
+        elif transform == "eps_greedy":
+            w = np.zeros_like(scores)
+            w[scores.argmax()] = 1.0 - self._eps
+            w += self._eps / self._n
+            return w
+        return np.ones_like(scores)
+
+    def task_scores(self) -> np.ndarray:
+        with self._lock:
+            return np.array(self._task_scores[:])
+
+    def metrics(self) -> dict:
+        with self._lock:
+            unseen = np.array(self._unseen_task_weights[:])
+            return {
+                "task_scores": np.array(self._task_scores[:]),
+                "task_score_steps": np.array(self._task_score_steps[:]),
+                "unseen_task_weights": unseen,
+                "task_staleness": np.array(self._task_staleness[:]),
+                "proportion_seen": self._proportion_seen(),
+            }
+
+    def _state_dict_locked(self) -> dict:
+        return {
+            "task_scores": list(self._task_scores[:]),
+            "task_score_steps": list(self._task_score_steps[:]),
+            "task_staleness": list(self._task_staleness[:]),
+            "unseen_task_weights": list(self._unseen_task_weights[:]),
+        }
+
+    def _load_state_dict_locked(self, state: dict) -> None:
+        self._task_scores[:] = [float(value) for value in state.get("task_scores", self._task_scores[:])]
+        self._task_score_steps[:] = [int(value) for value in state.get("task_score_steps", self._task_score_steps[:])]
+        self._task_staleness[:] = [float(value) for value in state.get("task_staleness", self._task_staleness[:])]
+        self._unseen_task_weights[:] = [
+            float(value) for value in state.get("unseen_task_weights", self._unseen_task_weights[:])
+        ]
+        self._recompute_weights()