Spaces:
Runtime error
Runtime error
| from typing import Iterable, Optional | |
| import numpy as np | |
| from .policies import ensemble_levo_thinking_policy | |
| class LevoThinkingEnsembleAgent: | |
| """ | |
| Ensemble Q-learning agent with LevoThinking policy. | |
| Q has shape (n_heads, n_states, n_actions). | |
| For action selection, we use: | |
| mu(a) = mean_h Q_h(s, a) | |
| var(a) = var_h Q_h(s, a) | |
| base_score(a) = mu(a) - lambda_var * var(a) | |
| score(a) = base_score(a) + HF(t, a) | |
| Each head is updated independently with standard Q-learning, optionally | |
| adding Gaussian noise to the TD-error for a subset of heads. | |
| """ | |
| def __init__( | |
| self, | |
| n_states: int, | |
| n_actions: int, | |
| n_heads: int = 5, | |
| gamma: float = 0.99, | |
| alpha: float = 0.1, | |
| A: float = 0.5, | |
| omega: float = 0.1, | |
| phase_offset: float = 0.0, | |
| lambda_var: float = 0.5, | |
| prior: Optional[np.ndarray] = None, | |
| tau: float = 1.0, | |
| name: str = "LevoThinking", | |
| ) -> None: | |
| self.n_states = n_states | |
| self.n_actions = n_actions | |
| self.n_heads = n_heads | |
| self.gamma = gamma | |
| self.alpha = alpha | |
| self.A = A | |
| self.omega = omega | |
| self.phase_offset = phase_offset | |
| self.lambda_var = lambda_var | |
| self.prior = prior | |
| self.tau = tau | |
| self.name = name | |
| self.Q = np.zeros((n_heads, n_states, n_actions), dtype=float) | |
| self.t = 0 # global step counter for the HF modulation | |
| def select_action(self, state: int, rng: np.random.Generator) -> int: | |
| probs = ensemble_levo_thinking_policy( | |
| self.Q, | |
| state, | |
| t=self.t, | |
| A=self.A, | |
| omega=self.omega, | |
| phase_offset=self.phase_offset, | |
| lambda_var=self.lambda_var, | |
| prior=self.prior, | |
| tau=self.tau, | |
| ) | |
| self.t += 1 | |
| return int(rng.choice(self.n_actions, p=probs)) | |
| def update( | |
| self, | |
| s: int, | |
| a: int, | |
| r: float, | |
| s_next: int, | |
| done: bool, | |
| noisy_heads: Optional[Iterable[int]] = None, | |
| noise_std: float = 0.0, | |
| rng: Optional[np.random.Generator] = None, | |
| ) -> None: | |
| """ | |
| Update each head with standard Q-learning. | |
| If `noisy_heads` is not None and `noise_std > 0`, then for any head h | |
| in `noisy_heads`, we add N(0, noise_std) to the TD-error before the | |
| update. This is used in Phase 3 to simulate corrupted learning signals. | |
| """ | |
| if noisy_heads is None: | |
| noisy_set = set() | |
| else: | |
| noisy_set = set(int(h) for h in noisy_heads) | |
| for h in range(self.n_heads): | |
| q_sa = self.Q[h, s, a] | |
| if done: | |
| target = r | |
| else: | |
| target = r + self.gamma * float(self.Q[h, s_next].max()) | |
| td = target - q_sa | |
| if h in noisy_set and noise_std > 0.0 and rng is not None: | |
| td += float(rng.normal(loc=0.0, scale=noise_std)) | |
| self.Q[h, s, a] = q_sa + self.alpha * td | |