The code below has a bug somewhere, but I don't know where, so I would like someone to take a look. The agent never learns — it just lands between -200 and -400 reward. I'm not sure if I'm mishandling the way parallel environments work. I created another version of the code that goes through a single step and works fine, but this one just doesn't. I asked some AIs like ChatGPT, Gemini, and Claude, but they didn't find any error. This code is analogous to a lab in a Coursera course that uses just NumPy to calculate the gradients and processes one environment at a time. But I tried to rewrite it using PyTorch and multiple environments at each step, so it can learn faster. Thanks in advance import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F import gymnasium as gym import numpy as np from tqdm import tqdm import copy torch.manual_seed(42) # Network to train the agent, should be initialize with a Dict with necessary parameters class ActionValueNetwork(nn.Module): def __init__(self, network_config: dict): super(ActionValueNetwork, self).__init__() self.state_dim = network_config["state_dim"] self.action_dim = network_config["action_dim"] self.hidden_dim = network_config["hidden_dim"] self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.network = nn.Sequential( nn.Linear(self.state_dim, self.hidden_dim), nn.ReLU(), nn.Linear(self.hidden_dim, self.action_dim) ) def forward(self, state): return self.network(state) class ReplayBuffer: def __init__(self, buffer_config:dict): self.max_size = buffer_config["buffer_size"] self.minibatch_size = buffer_config["minibatch_size"] self.seed = buffer_config["seed"] self.state_dim = buffer_config["state_dim"] # Initialize memory self.rand_generator = np.random.RandomState(self.seed) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Pre-allocate memory self.states = torch.zeros((self.max_size, self.state_dim), dtype=torch.float32, device=self.device) self.actions = torch.zeros(self.max_size, dtype=torch.long, device=self.device) self.rewards = torch.zeros(self.max_size, dtype=torch.float32, device=self.device) self.next_states = torch.zeros((self.max_size, self.state_dim), dtype=torch.float32, device=self.device) self.dones = torch.zeros(self.max_size, dtype=torch.bool, device=self.device) self.count = 0 self.current = 0 def append(self, state, action, reward, done, next_state=None): """Note: This method cannot be used in the same loop with append_batch""" # Convert to tensors if they aren't already if not isinstance(state, torch.Tensor): state = torch.tensor(state, dtype=torch.float32, device=self.device) if not isinstance(next_state, torch.Tensor): next_state = torch.tensor(next_state, dtype=torch.float32, device=self.device) if next_state is not None else torch.zeros_like(state, device=self.device) # Store experience self.states[self.current] = state.squeeze() self.actions[self.current] = torch.tensor(action, dtype=torch.long, device=self.device) self.rewards[self.current] = torch.tensor(reward, dtype=torch.float32, device=self.device) self.next_states[self.current] = next_state.squeeze() self.dones[self.current] = torch.tensor(done, dtype=torch.bool, device=self.device) self.count = max(self.count, self.current + 1) self.current = (self.current + 1) % self.max_size def append_batch(self, states, actions, rewards, dones, next_states=None): """Note: This method cannot be used in the same loop with append""" # Convert to tensor if they aren't already if not isinstance(states, torch.Tensor): states = torch.tensor(states, dtype=torch.float32, device=self.device) if not isinstance(next_states, torch.Tensor): next_states = torch.tensor(next_states, dtype=torch.float32, device=self.device) if next_states is not None else torch.zeros_like(states, device=self.device) env_p_size = states.shape[0] self.states[self.current:self.current + env_p_size] = states self.actions[self.current:self.current + env_p_size] = torch.tensor(actions, dtype=torch.long, device=self.device) self.rewards[self.current:self.current + env_p_size] = torch.tensor(rewards, dtype=torch.float32, device=self.device) self.next_states[self.current:self.current + env_p_size] = next_states self.dones[self.current:self.current + env_p_size] = torch.tensor(dones, dtype=torch.bool, device=self.device) self.count = max(self.count, self.current + env_p_size) self.current = (self.current + env_p_size) % self.max_size def sample(self): if self.count < self.minibatch_size: return None idxs = self.rand_generator.choice(self.count, size=self.minibatch_size) return ( self.states[idxs].to(self.device), self.actions[idxs].to(self.device), self.rewards[idxs].to(self.device), self.next_states[idxs].to(self.device), self.dones[idxs].to(self.device) ) @property def size(self): return self.count def get_td_error(states, next_states, actions, rewards, dones, gamma: float, policy: ActionValueNetwork, current_policy: ActionValueNetwork, tau: float = 1.0): # Get Q-values for current states from the main network q_values = policy(states).to(policy.device) batch_indices = torch.arange(states.shape[0], device=policy.device) # Ensure indices are on the correct device q_to_update = q_values[batch_indices, actions] # Get Q-values for next states from the main network (detached) with torch.no_grad(): # Ensure no gradients flow back from target calculation q_nexts = current_policy(next_states).to(policy.device) # Use policy network for next Q-values # --- Corrected Expected SARSA Target Calculation --- # Apply temperature scaling for softmax calculation (with stability trick) max_q_nexts, _ = torch.max(q_nexts, dim=1, keepdim=True) q_nexts_stable_for_softmax = q_nexts - max_q_nexts # For numerical stability in softmax q_nexts_tau = q_nexts_stable_for_softmax / tau # Calculate softmax probabilities based on temperature-scaled values probs = F.softmax(q_nexts_tau, dim=1) # Calculate the expected next state value using the ORIGINAL q_nexts values # E[Q(s', a')] = sum( softmax(Q(s',.)/tau) * Q(s',.) ) exp_sarsa = torch.sum(probs * q_nexts, dim=1) # Use original q_nexts here! # ------------------------------------------------- # Calculate TD target # Ensure dones is float for multiplication td_target = rewards + gamma * exp_sarsa * (1 - dones.float()) # Calculate TD error td_error = td_target - q_to_update #print(f"TD target: {td_target}") #print(f"Q to update: {q_to_update}") #print(f"TD error: {td_error}") return td_error def optimize_network_clipped(states, actions, rewards, next_states, dones, gamma, optimizer, policy, current_policy, tau): # Calculate the TD error td_error = get_td_error(states, next_states, actions, rewards, dones, gamma, policy, current_policy, tau) loss = torch.mean(torch.square(td_error)) optimizer.zero_grad() loss.backward() # Calculate clipped gradient torch.nn.utils.clip_grad_norm_(policy.parameters(), max_norm=10) optimizer.step() #print(f"Loss: {loss.item()}") return loss.item() class AgentParallel: def __init__(self, agent_config: dict): self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.num_envs = agent_config["num_envs"] self.replay_buffer = ReplayBuffer(agent_config["buffer_config"]) self.network = ActionValueNetwork(agent_config["network_config"]).to(self.device) self.optimizer = optim.Adam(self.network.parameters(), lr=agent_config["learning_rate"]) self.gamma = agent_config["gamma"] self.num_replay_updates = agent_config["num_replay_updates"] self.minibatch_size = self.replay_buffer.minibatch_size self.tau = agent_config["tau"] self.last_states = None self.last_actions = None self.sum_rewards = np.zeros(self.num_envs) self.episode_steps = np.zeros(self.num_envs, dtype=int) self.episodes_completed = 0 self.sum_rewards_history = [] self.update_network = 0 self.update_network_interval = agent_config["network_up_i"] def policy(self, states): if not isinstance(states, torch.Tensor): states = torch.tensor(states, dtype=torch.float32, device=self.device) with torch.no_grad(): action_values_b = self.network(states) max_val_b, _ = torch.max(action_values_b, dim=1, keepdim=True) actions_values_stable_b = action_values_b - max_val_b actions_values_temp_b = actions_values_stable_b / self.tau probs_b = F.softmax(actions_values_temp_b, dim=1) actions_distribution_b = torch.distributions.Categorical(probs=probs_b) actions_b = actions_distribution_b.sample() return actions_b.cpu().numpy() def agent_start(self, states): # Select an action for each state using the policy self.last_states = torch.tensor(states, dtype=torch.float32, device=self.device) self.last_actions = self.policy(self.last_states) # Rest count just for envs that start a new episode self.num_rewards = np.zeros(self.num_envs) self.episode_steps = np.zeros(self.num_envs, dtype=int) return self.last_actions def agent_step(self, rewards, states, dones): # Update sum reward and episode steps self.sum_rewards += rewards self.episode_steps += 1 # Make tensor states = torch.tensor(states, dtype=torch.float32, device=self.device) rewards = torch.tensor(rewards, dtype=torch.float32, device=self.device) dones = torch.tensor(dones, dtype=torch.bool, device=self.device) # Select an action actions = self.policy(states) # Store batch in replay buffer self.replay_buffer.append_batch(self.last_states, torch.tensor(self.last_actions, dtype=torch.long, device=self.device), rewards, dones, states) if self.replay_buffer.size > self.minibatch_size: if self.update_network == self.update_network_interval: self.update_network = 0 current_policy = copy.deepcopy(self.network) for _ in range(self.num_replay_updates): batch = self.replay_buffer.sample() if batch is not None: mb_states, mb_actions, mb_rewards, mb_next_states, mb_dones = batch loss = optimize_network_clipped(mb_states, mb_actions, mb_rewards, mb_next_states, mb_dones, self.gamma, self.optimizer, self.network, current_policy, self.tau) # Episode dones register for i, done in enumerate(dones): if done: self.episodes_completed += 1 self.sum_rewards_history.append(self.sum_rewards) # Reset rewards for enviroments completed self.sum_rewards = 0 self.episode_steps = 0 self.update_network += 1 self.last_states = states self.last_actions = actions return actions # Experiment num_envs = 8 env = gym.make_vec("LunarLander-v3", num_envs=num_envs, vectorization_mode='sync') tau = 0.1 minibatch_size = 32 gamma = 0.9 # Settings network_config = { "state_dim": env.single_observation_space.shape[0], "action_dim": int(env.single_action_space.n), "hidden_dim": 256 } buffer_config = { "buffer_size": 50000, "minibatch_size": minibatch_size, "seed": 42, "state_dim": env.single_observation_space.shape[0] } agent_config = { "network_config": network_config, "buffer_config": buffer_config, "learning_rate": 1e-3, "gamma": gamma, 'tau': tau, "num_replay_updates": 4, "minibatch_size": minibatch_size, "num_envs": num_envs } agent = AgentParallel(agent_config) target_episodes = 2000 print_interval = 100 observation, info = env.reset() action = agent.agent_start(observation) pbar = tqdm(total=target_episodes) last_episodes_completed = 0 while agent.episodes_completed < target_episodes: # Give a step in all enviroments observations, rewards, terminated_v, truncated_v, infos = env.step(action) # Combine teminated signals dones = np.logical_or(terminated_v, truncated_v) # Update Agent actions = agent.agent_step(rewards, observations, dones) if agent.episodes_completed > last_episodes_completed: pbar.update(agent.episodes_completed - last_episodes_completed) last_episodes_completed = agent.episodes_completed # Print metrics if agent.episodes_completed % print_interval == 0: recent_rewards = agent.sum_rewards_history[-print_interval:] avg_reward = sum(recent_rewards) / len(recent_rewards) print(f"\nEpisodes: {agent.episodes_completed}, Avg. Reward last ({print_interval}): {avg_reward:.2f}") pbar.close() env.close() Continue reading...
