| from __future__ import annotations |
|
|
| import random |
| from collections import Counter |
|
|
| from data.datasets.android_control_a11y import ACTION_TYPES, AndroidControlA11yDataset |
| from rl_harness.types import PolicyPrediction, Reward |
| from scripts.train_qwen35_history_actions import ACTION_TYPES_BY_GROUP, action_code, group_for_sample |
|
|
|
|
| class OfflineA11yBanditEnv: |
| """Single-step reward harness over held-out Android Control samples. |
| |
| This is the cheap preflight path for RL plumbing. It validates prompt assembly, |
| constrained decode, action parsing, label lookup, reward accounting, and logs. |
| It does not claim trajectory credit because the dataset cannot branch after a |
| model action. |
| """ |
|
|
| def __init__( |
| self, |
| shards: list[int], |
| history_len: int = 4, |
| min_history: int = 4, |
| max_raw_tree_chars: int = 5000, |
| quota: str = "tp=5,k=1,scroll=3,wait=2,system=1", |
| seed: int = 1, |
| ): |
| self.rng = random.Random(seed) |
| random.seed(seed) |
| excluded = set(range(20)) - set(shards) |
| self.ds = AndroidControlA11yDataset(history_len=history_len, start_shard=min(shards), excluded_shards=excluded) |
| self.min_history = min_history |
| self.max_raw_tree_chars = max_raw_tree_chars |
| self.quota = self._parse_quota(quota) |
| self.targets = [] |
| for group, count in self.quota.items(): |
| self.targets.extend([group] * count) |
| self._target_i = 0 |
|
|
| @staticmethod |
| def _parse_quota(text: str) -> Counter: |
| out = Counter() |
| for part in text.split(","): |
| key, value = part.split("=", 1) |
| out[key.strip()] = int(value) |
| return out |
|
|
| def sample(self) -> dict: |
| while True: |
| if self._target_i % max(len(self.targets), 1) == 0: |
| self.rng.shuffle(self.targets) |
| wanted = self.targets[self._target_i % len(self.targets)] |
| self._target_i += 1 |
| sample = self.ds.sample_random(allowed_action_types=ACTION_TYPES_BY_GROUP[wanted]) |
| if len(sample.get("history", [])) < self.min_history: |
| continue |
| if self.max_raw_tree_chars and len(sample.get("compressed_tree", "")) > self.max_raw_tree_chars: |
| continue |
| if group_for_sample(sample) != wanted: |
| continue |
| return sample |
|
|
| @staticmethod |
| def _score_strings(gt: str, pred: str, group: str) -> tuple[bool, bool, str]: |
| gt = gt.strip() |
| pred = pred.strip() |
| gt_action = gt.split(maxsplit=1)[0] |
| pred_action = pred.split(maxsplit=1)[0] if pred else "" |
| action_ok = pred_action == gt_action |
| if not action_ok: |
| return False, action_ok, "wrong_action" |
|
|
| if group == "k": |
| if "<eos>" not in pred: |
| return False, action_ok, "missing_eos" |
| pred_text = pred.split("<eos>", 1)[0].strip().lower() |
| gt_text = gt.lower() |
| return pred_text == gt_text, action_ok, "wrong_value" |
|
|
| return pred == gt, action_ok, "wrong_value" |
|
|
| def reward(self, sample: dict, prediction: PolicyPrediction, label_map: dict[int, str]) -> Reward: |
| gt = action_code(sample, label_map) |
| if gt is None: |
| return Reward(0.0, True, "unscoreable") |
| pred = prediction.raw.strip() |
| group = group_for_sample(sample) |
| exact, action_ok, miss_reason = self._score_strings(gt, pred, group) |
|
|
| if exact: |
| value = 1.0 |
| reason = "exact" |
| elif action_ok: |
| value = 0.25 |
| reason = f"right_action_{miss_reason}" |
| else: |
| value = 0.0 |
| reason = miss_reason |
|
|
| return Reward( |
| value=value, |
| done=True, |
| reason=reason, |
| details={ |
| "gt": gt, |
| "pred": pred, |
| "group": group, |
| "act_type": ACTION_TYPES[sample["act_type_idx"]], |
| "action_ok": action_ok, |
| "exact": exact, |
| }, |
| ) |
|
|
