Spaces:
Runtime error
Runtime error
File size: 3,222 Bytes
ccba775 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | 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
|