Spaces:
Runtime error
Runtime error
| from typing import Optional | |
| import numpy as np | |
| from .policies import hf_levo_policy_tabular | |
| class LevoQTabularAgent: | |
| """ | |
| Tabular Q-learning agent with an HF-Levo policy. | |
| The Q-update is standard: | |
| Q(s, a) <- Q(s, a) + alpha * (r + gamma max_a' Q(s', a') - Q(s, a)) | |
| The policy is: | |
| score(a) = Q(s, a) + A cos(omega * t + phi_a + phase_offset) | |
| pi(a|s) = softmax(score / tau) | |
| """ | |
| def __init__( | |
| self, | |
| n_states: int, | |
| n_actions: int, | |
| gamma: float = 0.99, | |
| alpha: float = 0.1, | |
| A: float = 0.5, | |
| omega: float = 0.1, | |
| phase_offset: float = 0.0, | |
| ent_weight: float = 0.0, | |
| prior: Optional[np.ndarray] = None, | |
| tau: float = 1.0, | |
| name: str = "LevoQ", | |
| ) -> None: | |
| self.n_states = n_states | |
| self.n_actions = n_actions | |
| self.gamma = gamma | |
| self.alpha = alpha | |
| self.A = A | |
| self.omega = omega | |
| self.phase_offset = phase_offset | |
| self.ent_weight = ent_weight | |
| self.prior = prior | |
| self.tau = tau | |
| self.name = name | |
| self.Q = np.zeros((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 = hf_levo_policy_tabular( | |
| self.Q, | |
| state, | |
| t=self.t, | |
| A=self.A, | |
| omega=self.omega, | |
| phase_offset=self.phase_offset, | |
| ent_weight=self.ent_weight, | |
| 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) -> None: | |
| q_sa = self.Q[s, a] | |
| if done: | |
| target = r | |
| else: | |
| target = r + self.gamma * float(self.Q[s_next].max()) | |
| td = target - q_sa | |
| self.Q[s, a] = q_sa + self.alpha * td | |