""" Tabular Q-Learning RL agent. Uses state hashing and epsilon-greedy exploration with action masking. Designed as a lightweight RL agent that can learn effective strategies without requiring deep learning frameworks. """ import random import json import hashlib from typing import Dict, Tuple, Optional from collections import defaultdict class RLAgent: """Tabular Q-learning agent with epsilon-greedy exploration. Features: - State hashing for tabular lookup. - Action masking (only considers valid actions). - Epsilon decay for explore → exploit transition. - Learning rate and discount factor tuning. """ def __init__( self, learning_rate: float = 0.1, discount_factor: float = 0.95, epsilon: float = 1.0, epsilon_min: float = 0.05, epsilon_decay: float = 0.995, seed: Optional[int] = None, ): """Initialize the Q-learning agent. Args: learning_rate: Alpha for Q-value updates. discount_factor: Gamma for future reward discounting. epsilon: Initial exploration rate. epsilon_min: Minimum exploration rate. epsilon_decay: Multiplicative decay per episode. seed: Random seed. """ self.lr = learning_rate self.gamma = discount_factor self.epsilon = epsilon self.epsilon_min = epsilon_min self.epsilon_decay = epsilon_decay # Q-table: state_hash → {action_key → Q-value} self.q_table: Dict[str, Dict[str, float]] = defaultdict( lambda: defaultdict(float) ) # Experience tracking self.last_state_hash: Optional[str] = None self.last_action_key: Optional[str] = None if seed is not None: random.seed(seed) def _hash_state(self, state: Dict) -> str: """Create a compact hash of the state for table lookup. We hash key features rather than the full state for generalization: - Current time - Number of pending/completed/missed tasks by priority - Number of unreplied messages by urgency """ features = { "time": state.get("time", ""), "pending_high": sum( 1 for t in state.get("tasks", []) if t["status"] == "pending" and t["priority"] == "high" ), "pending_med": sum( 1 for t in state.get("tasks", []) if t["status"] == "pending" and t["priority"] == "medium" ), "pending_low": sum( 1 for t in state.get("tasks", []) if t["status"] == "pending" and t["priority"] == "low" ), "completed": sum( 1 for t in state.get("tasks", []) if t["status"] == "completed" ), "missed": sum( 1 for t in state.get("tasks", []) if t["status"] == "missed" ), "urgent_unreplied": sum( 1 for m in state.get("inbox", []) if m.get("urgency") == "high" and not m.get("replied", False) ), "normal_unreplied": sum( 1 for m in state.get("inbox", []) if m.get("urgency") != "high" and not m.get("replied", False) ), } feature_str = json.dumps(features, sort_keys=True) return hashlib.md5(feature_str.encode()).hexdigest()[:12] def _action_key(self, action: Tuple[str, int]) -> str: """Convert action tuple to a string key for Q-table lookup.""" return f"{action[0]}:{action[1]}" def _parse_action_key(self, key: str) -> Tuple[str, int]: """Convert action key back to tuple.""" parts = key.split(":") return (parts[0], int(parts[1])) def act(self, state: Dict) -> Tuple[str, int]: """Choose an action using epsilon-greedy policy with action masking. Args: state: Observation dict from the environment. Returns: (action_type, target_id) tuple. """ valid_actions = state.get("valid_actions", []) if not valid_actions: return ("defer_task", 0) state_hash = self._hash_state(state) # Epsilon-greedy exploration if random.random() < self.epsilon: action = random.choice(valid_actions) else: # Exploit: choose best Q-value among valid actions q_values = self.q_table[state_hash] best_action = None best_q = float("-inf") for va in valid_actions: ak = self._action_key(va) q = q_values[ak] if q > best_q: best_q = q best_action = va action = best_action if best_action else random.choice(valid_actions) # Store for learning self.last_state_hash = state_hash self.last_action_key = self._action_key(action) return action def learn( self, reward: float, next_state: Dict, done: bool, ): """Update Q-values using the Q-learning update rule. Q(s,a) ← Q(s,a) + α[r + γ·max_a' Q(s',a') - Q(s,a)] Args: reward: Reward received from the last action. next_state: New observation after the action. done: Whether the episode ended. """ if self.last_state_hash is None or self.last_action_key is None: return current_q = self.q_table[self.last_state_hash][self.last_action_key] if done: target = reward else: # Compute max Q-value for next state (among valid actions) next_hash = self._hash_state(next_state) next_valid = next_state.get("valid_actions", []) if next_valid: max_next_q = max( self.q_table[next_hash][self._action_key(a)] for a in next_valid ) else: max_next_q = 0.0 target = reward + self.gamma * max_next_q # Q-learning update self.q_table[self.last_state_hash][self.last_action_key] = ( current_q + self.lr * (target - current_q) ) def decay_epsilon(self): """Decay exploration rate after each episode.""" self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay) def get_q_table_size(self) -> int: """Return the number of unique states seen.""" return len(self.q_table) def get_stats(self) -> Dict: """Return agent statistics.""" total_entries = sum(len(v) for v in self.q_table.values()) return { "q_table_states": len(self.q_table), "q_table_entries": total_entries, "epsilon": round(self.epsilon, 4), } def __repr__(self): return ( f"RLAgent(lr={self.lr}, gamma={self.gamma}, " f"epsilon={self.epsilon:.3f}, states={len(self.q_table)})" )