support task selector

Author: chenyushuo

trinity/buffer/buffer.py

@@ -20,14 +20,18 @@ def get_buffer_reader(storage_config: StorageConfig, buffer_config: BufferConfig
         from trinity.buffer.reader.file_reader import (
             ExperienceFileReader,
             TaskFileReader,
+            TaskFileReaderWithSelector,
         )
 
         schema_type = storage_config.schema_type
         if schema_type:
             # only trainer input has schema type
             return ExperienceFileReader(storage_config, buffer_config)
         else:
-            return TaskFileReader(storage_config, buffer_config)
+            if storage_config.task_selector:
+                return TaskFileReaderWithSelector(storage_config, buffer_config)
+            else:
+                return TaskFileReader(storage_config, buffer_config)
     else:
         raise ValueError(f"{storage_config.storage_type} not supported.")
 

trinity/buffer/buffer_reader.py

@@ -1,6 +1,6 @@
 """Reader of the buffer."""
 from abc import ABC, abstractmethod
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 
 class BufferReader(ABC):
@@ -13,3 +13,14 @@ def read(self, batch_size: Optional[int] = None) -> List:
     @abstractmethod
     async def read_async(self, batch_size: Optional[int] = None) -> List:
         """Read from buffer asynchronously."""
+
+    @property
+    @abstractmethod
+    def index(self) -> int:
+        """Get the current index."""
+
+    def state_dict(self) -> Dict:
+        return {}
+
+    def load_state_dict(self, state_dict: Dict) -> None:
+        pass

trinity/buffer/operators/mappers/pass_rate_calculator.py

@@ -0,0 +1,25 @@
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+import numpy as np
+from sqlalchemy import Tuple
+
+from trinity.buffer.operators.experience_operator import (
+    EXPERIENCE_OPERATORS,
+    ExperienceOperator,
+)
+from trinity.buffer.task_scheduler import TASKSET_SCHEDULE_METRIC
+from trinity.common.experience import Experience
+
+
+@EXPERIENCE_OPERATORS.register_module("pass_rate_calculator")
+class PassRateCalculator(ExperienceOperator):
+    def __init__(self, reward_shaping_configs: Optional[List[Dict]] = None):
+        self.reward_shaping_configs = reward_shaping_configs
+
+    def process(self, exps: List[Experience]) -> Tuple[List[Experience], Dict]:
+        raw_metric = defaultdict(list)
+        for exp in exps:
+            raw_metric[exp.task_index].append(exp.reward)
+        metric = {task_index: np.mean(rewards) for task_index, rewards in raw_metric.items()}
+        return exps, {TASKSET_SCHEDULE_METRIC: metric}

trinity/buffer/reader/diff_based_selector.py

@@ -0,0 +1,150 @@
+from typing import List
+
+import numpy as np
+import torch
+
+from trinity.common.config import DataSelectorConfig, StorageConfig
+
+from .diff_estimator import InterpolationBetaPREstimator
+
+
+def build_diff_estimator(dataset, config: DataSelectorConfig):
+    print(f"[DEBUG]: {config=}")
+    feature_keys = config.feature_keys
+    features = np.concat([np.array(list(dataset[k]))[:, None] for k in feature_keys], axis=1)
+    print(f"[DEBUG]: {features.shape=}")
+    print(f"[DEBUG]: {features[:5]=}")
+    adaptive_rho = hasattr(config, "adaptive_rho") and config.adaptive_rho
+    return InterpolationBetaPREstimator(
+        features=features, m=config.m, lamb=config.lamb, rho=config.rho, adaptive_rho=adaptive_rho
+    )
+
+
+class BaseSelector:
+    def __init__(self, data_source, config: DataSelectorConfig):
+        self.data_source = data_source
+        self.config = config
+
+    def get_indices(self, batch_size: int, return_extra_info: bool = False):
+        raise NotImplementedError
+
+    def update(self, indices: List[int], values: List[float]):
+        raise NotImplementedError
+
+
+class RandomSelector(BaseSelector):
+    def __init__(self, data_source, config: DataSelectorConfig):
+        super().__init__(data_source, config)
+        self.n = len(data_source)
+        print(f"[DEBUG]: RandomSelector-{self.n=}")
+
+    def get_indices(self, batch_size, return_extra_info=False):
+        selected_indices = torch.from_numpy(np.random.permutation(self.n)[:batch_size])
+        print(f"[DEBUG]: RandomSelector-{selected_indices=}")
+        if return_extra_info:
+            return selected_indices, {}
+        else:
+            return selected_indices
+
+    def update(self, *args, **kwargs):
+        pass
+
+
+class OfflineEasy2HardSelector(BaseSelector):
+    def __init__(self, data_source, config: DataSelectorConfig):
+        super().__init__(data_source, config)
+
+        feature_keys = config.feature_keys
+        self.features = np.concat(
+            [np.array(list(data_source[k]))[:, None] for k in feature_keys], axis=1
+        )
+        features_with_index = [list(self.features[i]) + [i] for i in range(len(self.features))]
+        features_with_index = sorted(features_with_index)[::-1]
+        print(f"[DEBUG]: OfflineEasy2HardSelector, sorted {features_with_index[:20]}")
+        self.sorted_index = np.array([i[2] for i in features_with_index])
+
+        self.n = len(data_source)
+        self.current_position = 0
+
+    def update(self, *args, **kwargs) -> None:
+        pass
+
+    def get_indices(self, batch_size, return_extra_info=False):
+        if self.current_position + batch_size > self.n:
+            new_position = self.current_position + batch_size - self.n
+            selected_indices = np.concatenate(
+                [self.sorted_index[self.current_position :], self.sorted_index[:new_position]]
+            )
+        else:
+            new_position = self.current_position + batch_size
+            selected_indices = self.sorted_index[self.current_position : new_position]
+        self.current_position = new_position
+        if not return_extra_info:
+            return selected_indices
+        else:
+            extra_info = {
+                "indices": selected_indices.tolist(),
+                "feat1": self.features[selected_indices, 0].tolist(),
+                "feat2": self.features[selected_indices, 1].tolist(),
+            }
+            return selected_indices, extra_info
+
+
+class DiffBasedSelector(BaseSelector):
+    def __init__(self, data_source, config: DataSelectorConfig) -> None:
+        super().__init__(data_source, config)
+        self.diff_estimator = build_diff_estimator(data_source, config)
+
+    def update(self, indices: List[int], values: List[float]) -> None:
+        self.diff_estimator.update(indices, values)
+
+    def get_scores(self) -> List[float]:
+        predicted_pr = self.diff_estimator.predict_pr(do_sample=self.config.do_sample)
+        scores = -np.abs(self.config.target_reward - predicted_pr)
+        return scores
+
+    def get_indices(self, batch_size, return_extra_info=False):
+        sampling_scores = self.get_scores()
+        sampling_scores = torch.from_numpy(sampling_scores)
+        if self.config.tau == 0:
+            selected_indices = torch.topk(sampling_scores, batch_size).indices
+        else:
+            sampling_logits = sampling_scores / self.config.tau
+            sampling_logits -= sampling_logits.max()
+            sampling_probabilities = torch.softmax(sampling_logits, dim=0)
+            selected_indices = torch.multinomial(
+                sampling_probabilities, batch_size, replacement=False
+            )
+        print(f"[DEBUG]: {selected_indices=}")
+        print(f"[DEBUG]: {sampling_scores=}")
+        print(f"[DEBUG]: {sampling_scores[selected_indices]=}")
+
+        if return_extra_info:
+            selected_indices_list = selected_indices.tolist()
+            alphas = self.diff_estimator.alphas[selected_indices_list]
+            betas = self.diff_estimator.betas[selected_indices_list]
+            point_est = alphas / (alphas + betas)
+            extra_info = {
+                "indices": selected_indices_list,
+                "scores": sampling_scores[selected_indices].tolist(),
+                "alphas": alphas.tolist(),
+                "betas": betas.tolist(),
+                "point": point_est.tolist(),
+            }
+            return selected_indices, extra_info
+        else:
+            return selected_indices
+
+
+def build_selector(dataset, config: StorageConfig) -> BaseSelector:
+    selector_config = config.task_selector
+    assert selector_config is not None
+    selector_type = selector_config.selector_type
+    if selector_type == "random":
+        return RandomSelector(dataset, selector_config)
+    elif selector_type == "diff":
+        return DiffBasedSelector(dataset, selector_config)
+    elif selector_type == "offline":
+        return OfflineEasy2HardSelector(dataset, selector_config)
+    else:
+        raise ValueError(f"Unknown selector type: {selector_type}")

trinity/buffer/reader/diff_estimator.py

@@ -0,0 +1,112 @@
+from typing import List
+
+import numpy as np
+
+
+class BaseBetaPREstimator:
+    n: int
+    m: int
+    lamb: float
+    rho: float
+    alphas: np.ndarray
+    betas: np.ndarray
+
+    def __init__(self, n: int, m: int = 16, lamb: float = 0.2, rho: float = 0.2):
+        """
+        alpha_{t+1} = (1 - lamb) * alpha_t + (1 - rho) * bar{s} + rho * tilde{s}
+        beta_{t+1} = (1 - lamb) beta_t + (1 - rho) * bar{f} + rho * tilde{f}
+        :param n: number of tasks
+        :param m: repeat times per tasks
+        :param lamb: discount factor of historical estimation
+        :param rho: weight of pseudo counts
+        """
+        self.n = n
+        self.m = m
+        self.lamb = lamb
+        self.rho = rho
+        self.alphas = np.ones(n, dtype=float)
+        self.betas = np.ones(n, dtype=float)
+        print(
+            f"[DEBUG] {self.n=}, {self.m=}, {self.lamb=}, {self.rho=}, {self.alphas=}, {self.betas=}"
+        )
+
+    def set(self, alphas, betas):
+        self.alphas = alphas
+        self.betas = betas
+
+    def _update(self, s_bar, f_bar, p_tilde):
+        self.alphas = (
+            (1 - self.lamb) * self.alphas
+            + self.lamb
+            + (1 - self.rho) * s_bar
+            + self.rho * p_tilde * self.m
+        )
+        self.betas = (
+            (1 - self.lamb) * self.betas
+            + self.lamb
+            + (1 - self.rho) * f_bar
+            + self.rho * (1 - p_tilde) * self.m
+        )
+
+    def update(self, ref_indices: List[int], ref_pass_rates: List[float]):
+        raise NotImplementedError
+
+    def predict_pr(self, indices=None, do_sample=False):
+        if indices is None:
+            indices = np.arange(self.n)
+        if not do_sample:
+            return self.alphas[indices] / (self.alphas[indices] + self.betas[indices])
+        else:
+            return np.random.beta(self.alphas[indices], self.betas[indices])
+
+    def equivalent_count(self, indices=None):
+        if indices is None:
+            indices = np.arange(self.n)
+        return self.alphas[indices] + self.betas[indices]
+
+
+class InterpolationBetaPREstimator(BaseBetaPREstimator):
+    def __init__(
+        self,
+        features: np.ndarray,
+        m: int,
+        lamb,
+        rho,
+        cap_coef_update_discount=0.9,
+        adaptive_rho=False,
+    ):
+        super(InterpolationBetaPREstimator, self).__init__(len(features), m, lamb, rho)
+        self.features = features  # [D, 2]
+        self.cap_coef = None
+        self.cap_coef_update_discount = cap_coef_update_discount
+        self.adaptive_rho = adaptive_rho
+
+    def update(self, ref_indices: List[int], ref_pass_rates: List[float]):
+        ref_pass_rate = np.mean(ref_pass_rates)
+        ref_anchor_pass_rates = np.mean(self.features[ref_indices], axis=0)
+        cap_estimate = (ref_pass_rate - ref_anchor_pass_rates[0]) / (
+            ref_anchor_pass_rates[1] - ref_anchor_pass_rates[0] + 1e-6
+        )
+        if self.cap_coef is None:
+            self.cap_coef = cap_estimate
+        else:
+            self.cap_coef = (
+                self.cap_coef_update_discount * self.cap_coef
+                + (1 - self.cap_coef_update_discount) * cap_estimate
+            )
+        s_bar = np.zeros(self.n, dtype=float)
+        s_bar[ref_indices] = np.array(ref_pass_rates) * self.m
+        f_bar = np.zeros(self.n, dtype=float)
+        f_bar[ref_indices] = (1 - np.array(ref_pass_rates)) * self.m
+        p_tilde = np.clip(
+            (self.features[:, 1] - self.features[:, 0]) * self.cap_coef + self.features[:, 0], 0, 1
+        )
+
+        predicted_pass_rates = p_tilde[ref_indices]
+        mean_abs_error = np.mean(np.abs(np.array(predicted_pass_rates) - np.array(ref_pass_rates)))
+        if self.adaptive_rho and mean_abs_error >= 0.25:
+            self.rho = self.rho * 0.5
+        print(f"[DEBUG]: {mean_abs_error=}, {self.rho=}")
+        p_tilde[ref_indices] = np.array(ref_pass_rates)
+
+        self._update(s_bar, f_bar, p_tilde)