Spaces:

tensorus
/

api

Sleeping

api / rl_agent.py

Upload 11 files

aa654a4 verified over 1 year ago

19.6 kB

# 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")

	# 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")