# rl_agent.py """ Implements a Deep Q-Network (DQN) Reinforcement Learning agent for Tensorus. The agent interacts with an environment, stores experiences (S, A, R, S', D) in TensorStorage, samples experiences, and trains its Q-network. Note on Experience Storage: - Large tensors (state, next_state) are stored individually in a 'rl_states' dataset. - Experience tuples containing IDs of states/next_states and scalar action/reward/done are stored as metadata in a placeholder tensor within the 'rl_experiences' dataset. - This approach balances tensor-native storage with manageable metadata, but sampling requires retrieving linked state tensors, which might be slow depending on storage backend. """ from typing import Any import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F import random import math import logging from typing import Tuple, Optional, Dict, Any # Import necessary Tensorus components from tensor_storage import TensorStorage from dummy_env import DummyEnv # Import our dummy environment # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # --- DQN Network Definition --- class DQN(nn.Module): """Simple MLP Q-Network.""" def __init__(self, n_observations: int, n_actions: int, hidden_size: int = 128): super(DQN, self).__init__() self.layer1 = nn.Linear(n_observations, hidden_size) self.layer2 = nn.Linear(hidden_size, hidden_size) self.layer3 = nn.Linear(hidden_size, n_actions) def forward(self, x: torch.Tensor) -> torch.Tensor: """Forward pass: returns Q-values for each action.""" # Ensure input is float if x.dtype != torch.float32: x = x.float() x = F.relu(self.layer1(x)) x = F.relu(self.layer2(x)) return self.layer3(x) # --- RL Agent Class --- class RLAgent: """DQN Agent interacting with TensorStorage.""" def __init__(self, tensor_storage: TensorStorage, state_dim: int, action_dim: int, hidden_size: int = 128, lr: float = 1e-4, gamma: float = 0.99, epsilon_start: float = 0.9, epsilon_end: float = 0.05, epsilon_decay: int = 10000, target_update_freq: int = 500, batch_size: int = 128, experience_dataset: str = "rl_experiences", state_dataset: str = "rl_states"): """ Initializes the RL Agent. Args: tensor_storage: Instance of TensorStorage. state_dim: Dimensionality of the environment state. action_dim: Number of discrete actions. hidden_size: Hidden layer size for the DQN. lr: Learning rate for the optimizer. gamma: Discount factor for future rewards. epsilon_*: Epsilon-greedy exploration parameters. target_update_freq: How often (in steps) to update the target network. batch_size: Number of experiences to sample for training. experience_dataset: Name of the dataset to store experience metadata. state_dataset: Name of the dataset to store state/next_state tensors. """ if not isinstance(tensor_storage, TensorStorage): raise TypeError("tensor_storage must be an instance of TensorStorage") self.tensor_storage = tensor_storage self.state_dim = state_dim self.action_dim = action_dim self.gamma = gamma self.batch_size = batch_size self.epsilon_start = epsilon_start self.epsilon_end = epsilon_end self.epsilon_decay = epsilon_decay self.target_update_freq = target_update_freq self.experience_dataset = experience_dataset self.state_dataset = state_dataset # Ensure datasets exist for ds_name in [self.experience_dataset, self.state_dataset]: try: self.tensor_storage.get_dataset(ds_name) except ValueError: logger.info(f"Dataset '{ds_name}' not found. Creating it.") self.tensor_storage.create_dataset(ds_name) # Device configuration (use GPU if available) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") logger.info(f"RL Agent using device: {self.device}") # Networks self.policy_net = DQN(state_dim, action_dim, hidden_size).to(self.device) self.target_net = DQN(state_dim, action_dim, hidden_size).to(self.device) self.target_net.load_state_dict(self.policy_net.state_dict()) self.target_net.eval() # Target network is only for inference # Optimizer self.optimizer = optim.AdamW(self.policy_net.parameters(), lr=lr, amsgrad=True) self.steps_done = 0 # Counter for epsilon decay and target updates def select_action(self, state: torch.Tensor) -> int: """Selects an action using epsilon-greedy strategy.""" sample = random.random() # Calculate current epsilon eps_threshold = self.epsilon_end + (self.epsilon_start - self.epsilon_end) * \ math.exp(-1. * self.steps_done / self.epsilon_decay) self.steps_done += 1 # Increment step counter here if sample > eps_threshold: # Exploit: choose the best action from the policy network with torch.no_grad(): # Ensure state is on the correct device and has batch dimension state = state.unsqueeze(0).to(self.device) if state.ndim == 1 else state.to(self.device) q_values = self.policy_net(state) action = q_values.max(1)[1].view(1, 1).item() # Get index of max Q value logger.debug(f"Exploiting: Q-Values={q_values.cpu().numpy()}, Chosen Action={action}") return action else: # Explore: choose a random action action = random.randrange(self.action_dim) logger.debug(f"Exploring: Chosen Action={action}") return action def store_experience(self, state: torch.Tensor, action: int, reward: float, next_state: Optional[torch.Tensor], done: bool) -> None: """Stores an experience tuple in TensorStorage.""" if state is None: logger.error("Cannot store experience with None state.") return # 1. Store state tensor state_id = self.tensor_storage.insert(self.state_dataset, state.cpu(), metadata={"component": "state"}) # 2. Store next_state tensor (if not None) next_state_id = None if next_state is not None: next_state_id = self.tensor_storage.insert(self.state_dataset, next_state.cpu(), metadata={"component": "next_state"}) # 3. Create experience metadata experience_metadata = { "state_id": state_id, "action": action, # Store action directly (assuming discrete & scalar) "reward": reward, # Store reward directly "next_state_id": next_state_id, # Can be None if done "done": int(done) # Store boolean as int } # 4. Store placeholder tensor with experience metadata # Using a small tensor as the primary data for this record placeholder_tensor = torch.tensor([1.0]) exp_record_id = self.tensor_storage.insert(self.experience_dataset, placeholder_tensor, experience_metadata) logger.debug(f"Stored experience record {exp_record_id}: state_id={state_id}, action={action}, reward={reward:.2f}, next_state_id={next_state_id}, done={done}") def optimize_model(self) -> None: """Performs one step of optimization on the policy network.""" # Check if enough samples are available in the experience dataset try: # A bit inefficient to get the full list just for the count, # TensorStorage could be enhanced with a count method. experience_count = len(self.tensor_storage.get_dataset(self.experience_dataset)) except ValueError: experience_count = 0 if experience_count < self.batch_size: logger.debug(f"Not enough experiences ({experience_count}/{self.batch_size}) to optimize yet.") return # 1. Sample experience metadata records try: sampled_metadata_records = self.tensor_storage.sample_dataset(self.experience_dataset, self.batch_size) except ValueError: logger.error(f"Could not sample from dataset {self.experience_dataset}") return except Exception as e: logger.error(f"Error sampling experiences: {e}", exc_info=True) return # 2. Retrieve actual state/next_state tensors based on IDs in metadata states = [] actions = [] rewards = [] next_states = [] dones = [] non_final_mask_list = [] # Keep track of which next_states are not None for record in sampled_metadata_records: meta = record['metadata'] state_id = meta.get('state_id') next_state_id = meta.get('next_state_id') action = meta.get('action') reward = meta.get('reward') done_flag = bool(meta.get('done', 1)) # Default to True if missing? Risky. Assume present. # Basic validation if state_id is None or action is None or reward is None: logger.warning(f"Skipping invalid sampled record: {meta}") continue # Retrieve state state_record = self.tensor_storage.get_tensor_by_id(self.state_dataset, state_id) if state_record is None: logger.warning(f"Could not find state tensor with ID {state_id} for experience {meta.get('record_id')}. Skipping sample.") continue states.append(state_record['tensor']) # Retrieve next state if it exists (i.e., not a terminal state) current_next_state = None if not done_flag and next_state_id: next_state_record = self.tensor_storage.get_tensor_by_id(self.state_dataset, next_state_id) if next_state_record: current_next_state = next_state_record['tensor'] else: # This shouldn't happen if storage is consistent, but handle it logger.warning(f"Could not find next_state tensor with ID {next_state_id} for non-terminal experience {meta.get('record_id')}. Treating as terminal.") done_flag = True # Treat as done if next state is missing next_states.append(current_next_state) # Will be None for terminal states non_final_mask_list.append(not done_flag) actions.append(torch.tensor([[action]], dtype=torch.long)) # Action needs to be [[action]] for gather() rewards.append(torch.tensor([reward], dtype=torch.float32)) dones.append(done_flag) # Keep Python bool for now, convert later if needed # If not enough valid samples were retrieved after lookup if not states: logger.warning("No valid samples retrieved after state lookup. Optimization step skipped.") return # 3. Batch the retrieved data # Filter out None values in next_states for target Q calculation non_final_next_states = torch.cat([ns for ns in next_states if ns is not None]).to(self.device) if any(non_final_mask_list) else None state_batch = torch.cat(states).to(self.device) action_batch = torch.cat(actions).to(self.device) reward_batch = torch.cat(rewards).to(self.device) non_final_mask = torch.tensor(non_final_mask_list, dtype=torch.bool, device=self.device) # 4. Compute Q(s_t, a) - the model computes Q(s_t), then we select the columns of actions taken # Ensure state_batch has correct dimensions if needed (e.g., B x C x H x W for images) # Our dummy env state is simple (B x 1) if state_batch.ndim == 1: # Ensure batch dimension exists state_batch = state_batch.unsqueeze(-1) state_action_values = self.policy_net(state_batch).gather(1, action_batch) # 5. Compute V(s_{t+1}) for all next states. # Expected values of actions for non_final_next_states are computed based # on the "older" target_net; selecting their best reward with max(1)[0]. next_state_values = torch.zeros(len(states), device=self.device) # Start with zeros for all if non_final_next_states is not None and non_final_next_states.numel() > 0: with torch.no_grad(): if non_final_next_states.ndim == 1: # Ensure batch dimension non_final_next_states = non_final_next_states.unsqueeze(-1) next_state_values[non_final_mask] = self.target_net(non_final_next_states).max(1)[0] # 6. Compute the expected Q values (Bellman equation) expected_state_action_values = (next_state_values * self.gamma) + reward_batch # 7. Compute loss criterion = nn.SmoothL1Loss() # Huber loss loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1)) # Ensure target has same shape # 8. Optimize the model self.optimizer.zero_grad() loss.backward() # In-place gradient clipping (prevents exploding gradients) torch.nn.utils.clip_grad_value_(self.policy_net.parameters(), 100) self.optimizer.step() logger.debug(f"Optimization step done. Loss: {loss.item():.4f}") # 9. Periodically update the target network if self.steps_done % self.target_update_freq == 0: self._update_target_network() def _update_target_network(self): """Copies weights from policy_net to target_net.""" logger.info(f"Updating target network at step {self.steps_done}") self.target_net.load_state_dict(self.policy_net.state_dict()) def train(self, env: DummyEnv, num_episodes: int): """Runs the training loop for a number of episodes.""" logger.info(f"--- Starting Training for {num_episodes} episodes ---") episode_rewards = [] for i_episode in range(num_episodes): state = env.reset() # state should be a tensor from env.reset() # state = torch.tensor(state, dtype=torch.float32, device=self.device).unsqueeze(0) # Add batch dim # Note: env.reset now returns a tensor, no need to convert here done = False current_episode_reward = 0 steps_in_episode = 0 while not done: # Select action action = self.select_action(state) # state is already a tensor # Perform action in environment next_state, reward, done, _ = env.step(action) # next_state is a tensor current_episode_reward += reward steps_in_episode += 1 # Store experience in TensorStorage # Note: store_experience handles moving tensors to CPU for storage if needed self.store_experience(state, action, reward, next_state if not done else None, done) # Move to the next state state = next_state # Perform one step of the optimization (on the policy network) self.optimize_model() # Check if env forces done (e.g. max steps reached in DummyEnv) if done: break # Exit loop if env says done episode_rewards.append(current_episode_reward) logger.info(f"Episode {i_episode+1}/{num_episodes} finished after {steps_in_episode} steps. Reward: {current_episode_reward:.2f}. Epsilon: {self.epsilon_end + (self.epsilon_start - self.epsilon_end) * math.exp(-1. * self.steps_done / self.epsilon_decay):.3f}") # Optional: Add plotting or saving logic here if (i_episode + 1) % 50 == 0: # Log average reward periodically avg_reward = sum(episode_rewards[-50:]) / len(episode_rewards[-50:]) logger.info(f" Average reward over last 50 episodes: {avg_reward:.2f}") logger.info("--- Training Finished ---") return episode_rewards # --- Example Usage --- if __name__ == "__main__": logger.info("--- Starting RL Agent Example ---") # 1. Setup TensorStorage storage = TensorStorage() # 2. Setup Dummy Environment env = DummyEnv(max_steps=100) # Episodes are max 100 steps # 3. Create the RL Agent agent = RLAgent( tensor_storage=storage, state_dim=env.state_dim, action_dim=env.action_dim, hidden_size=64, # Smaller network for simple env lr=5e-4, gamma=0.95, epsilon_start=0.95, epsilon_end=0.05, epsilon_decay=20000, # Slower decay for demonstration target_update_freq=200, # Update target net less frequently batch_size=64, # Smaller batch size experience_dataset="dummy_env_experiences", # Use specific names state_dataset="dummy_env_states" ) # 4. Run the training loop num_episodes_to_run = 200 # Adjust as needed rewards = agent.train(env, num_episodes=num_episodes_to_run) # 5. Optional: Check TensorStorage contents print("\n--- Checking TensorStorage contents (Sample) ---") try: exp_count = len(storage.get_dataset(agent.experience_dataset)) state_count = len(storage.get_dataset(agent.state_dataset)) print(f"Found {exp_count} experience records in '{agent.experience_dataset}'.") print(f"Found {state_count} state records in '{agent.state_dataset}'.") if exp_count > 0: print("\nExample experience record (metadata):") sample_exp = storage.sample_dataset(agent.experience_dataset, 1) if sample_exp: print(sample_exp[0]['metadata']) state_id = sample_exp[0]['metadata'].get('state_id') if state_id: state_rec = storage.get_tensor_by_id(agent.state_dataset, state_id) if state_rec: print(f" -> Corresponding state tensor (retrieved): {state_rec['tensor']}") except ValueError as e: print(f"Could not retrieve datasets: {e}") except Exception as e: print(f"An error occurred checking storage: {e}") logger.info("--- RL Agent Example Finished ---") # Optional: Plot rewards try: import matplotlib.pyplot as plt plt.figure(figsize=(10, 5)) plt.plot(rewards) # Calculate a simple moving average moving_avg = [sum(rewards[max(0, i-20):i+1])/len(rewards[max(0, i-20):i+1]) for i in range(len(rewards))] plt.plot(moving_avg, linestyle='--', label='Moving Avg (20 episodes)') plt.title("Episode Rewards over Time") plt.xlabel("Episode") plt.ylabel("Total Reward") plt.grid(True) plt.legend() print("\nPlotting rewards... Close the plot window to exit.") plt.show() except ImportError: print("\nMatplotlib not found. Skipping reward plot. Install with: pip install matplotlib")