Observation, Advantage and Return normalization for SAC and PPO

Observation_norm_SAC/Obser_norm_in_batch/sac_helpers_cnn_in_batch.py

@@ -0,0 +1,303 @@
+import numpy as np
+import torch as T
+import torch.nn as nn
+import torch.optim as optim
+from torch.distributions import Categorical
+
+class Agent:
+    def __init__(self, obs_space, action_space, hidden, gamma, lr, alpha, seed, batch_size, tau=0.005):
+        if seed is not None:
+            np.random.seed(seed)
+            T.manual_seed(seed)
+        
+        # Use GPU if available
+
+        self.device = T.device('cuda:0' if T.cuda.is_available() else 'cpu')
+        self.action_dim = int(getattr(action_space, "n", action_space.n))  # Use .n for Discrete
+        self.obs_shape = obs_space.shape
+        self.observeNorm = ObservationNorm()
+        self.gamma, self.tau, self.batch_size = gamma, tau, batch_size
+        # Make alpha learnable (adjust entropy based on reward magnitude)
+        self.target_entropy = -float(self.action_dim)
+        self.log_alpha = T.tensor(np.log(alpha), requires_grad=True, device=self.device)
+        self.alpha = np.exp(self.log_alpha.item()) 
+        self.alpha_opt = optim.Adam([self.log_alpha], lr=lr)
+
+        self.policy = CategoricalActor(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q1 = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q2 = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q1_target = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q2_target = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q1_target.load_state_dict(self.q1.state_dict())
+        self.q2_target.load_state_dict(self.q2.state_dict())
+
+        self.policy_opt = optim.Adam(self.policy.parameters(), lr=lr)
+        self.q1_opt = optim.Adam(self.q1.parameters(), lr=lr)
+        self.q2_opt = optim.Adam(self.q2.parameters(), lr=lr)
+        self.memory = Memory()
+
+    def choose_action(self, observation, eval_mode=False):
+        state = T.as_tensor(observation, dtype=T.float32, device=self.device)
+        state = self.observeNorm.normalize(state.unsqueeze(0))
+        state = state.squeeze(0)
+        with T.no_grad():
+            logits = self.policy(state.unsqueeze(0))
+            dist = Categorical(logits=logits)
+            if eval_mode:
+                action = logits.argmax(dim=-1)
+            else:
+                action = dist.sample()
+        return int(action.item())
+
+    def remember(self, state, action, reward, done, next_state):
+
+        state = T.as_tensor(state, dtype=T.float32, device=self.device)
+        if state.dim() == 3:  # [C,H,W]
+            state = state.unsqueeze(0)  # [1,C,H,W]
+
+        state = self.observeNorm.normalize(state)
+        state = state.squeeze(0)
+        # next_state also needs normalization
+
+        next_state = T.as_tensor(next_state, dtype=T.float32, device=self.device)
+        if next_state.dim() == 3:  # [C,H,W]
+            next_state = next_state.unsqueeze(0)  # [1,C,H,W]
+
+        next_state = self.observeNorm.normalize(next_state)
+        next_state = next_state.squeeze(0)
+
+        self.memory.store(state, action, reward, done, next_state)
+
+    def vanilla_sac_update(self):
+        if len(self.memory.states) < self.batch_size:
+            return 0.0
+        
+        # Mini-batch sampling
+        idxs = np.random.choice(len(self.memory.states), self.batch_size, replace=False)
+        states = T.as_tensor(np.array([self.memory.states[i] for i in idxs]), dtype=T.float32, device=self.device)
+        actions = T.as_tensor(np.array([self.memory.actions[i] for i in idxs]), dtype=T.int64, device=self.device).unsqueeze(-1)
+        rewards = T.as_tensor(np.array([self.memory.rewards[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        dones = T.as_tensor(np.array([self.memory.dones[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        next_states = T.as_tensor(np.array([self.memory.next_states[i] for i in idxs]), dtype=T.float32, device=self.device)
+
+        # Critic update, Soft Q-Learning Objective: to ensure high-entropy actions for exploration 
+        with T.no_grad():
+            next_logits = self.policy(next_states)
+            next_dist = Categorical(logits=next_logits)
+            next_probs = next_dist.probs
+            next_log_probs = next_dist.logits - T.logsumexp(next_dist.logits, dim=-1, keepdim=True)
+            q1_next = self.q1_target(next_states)
+            q2_next = self.q2_target(next_states)
+            # Soft Policy Evaluation 
+            min_q_next = T.min(q1_next, q2_next)
+            next_value = (next_probs * (min_q_next - self.alpha * next_log_probs)).sum(dim=-1, keepdim=True)
+            target = rewards + self.gamma * (1 - dones) * next_value
+
+        q1 = self.q1(states).gather(1, actions)
+        q2 = self.q2(states).gather(1, actions)
+
+        # Losses of both Q-functions
+        q1_loss = nn.MSELoss()(q1, target)
+        q2_loss = nn.MSELoss()(q2, target)
+
+        self.q1_opt.zero_grad()
+        q1_loss.backward()
+        self.q1_opt.step()
+        self.q2_opt.zero_grad()
+        q2_loss.backward()
+        self.q2_opt.step()
+
+        # Policy/Actor Objective
+        logits = self.policy(states)
+        dist = Categorical(logits=logits)
+        probs = dist.probs
+        log_probs = dist.logits - T.logsumexp(dist.logits, dim=-1, keepdim=True)
+        q1_policy = self.q1(states)
+        q2_policy = self.q2(states)
+        min_q_policy = T.min(q1_policy, q2_policy)
+        # Slightly different policy loss for discrete actions
+        policy_loss = (probs * (self.alpha * log_probs - min_q_policy)).sum(dim=-1).mean()
+
+        # Temperature to update Alpha
+        alpha_loss = -(self.log_alpha * (log_probs + self.target_entropy).detach()).mean()
+        self.alpha_opt.zero_grad()
+        alpha_loss.backward()
+        self.alpha_opt.step()
+        self.alpha = self.log_alpha.exp().item()  
+
+        self.policy_opt.zero_grad()
+        policy_loss.backward()
+        self.policy_opt.step()
+
+        # Target network update
+        for target_param, param in zip(self.q1_target.parameters(), self.q1.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+        for target_param, param in zip(self.q2_target.parameters(), self.q2.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+
+        return policy_loss.item()
+    
+    def update_reward_gradient_clipping(self):
+        if len(self.memory.states) < self.batch_size:
+            return 0.0
+        
+        # Mini-batch sampling
+        idxs = np.random.choice(len(self.memory.states), self.batch_size, replace=False)
+        states = T.as_tensor(np.array([self.memory.states[i] for i in idxs]), dtype=T.float32, device=self.device)
+        actions = T.as_tensor(np.array([self.memory.actions[i] for i in idxs]), dtype=T.int64, device=self.device).unsqueeze(-1)
+        rewards = T.as_tensor(np.array([self.memory.rewards[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        dones = T.as_tensor(np.array([self.memory.dones[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        next_states = T.as_tensor(np.array([self.memory.next_states[i] for i in idxs]), dtype=T.float32, device=self.device)
+        
+        """
+        # Min-max normalization and tanh scaling to [-1, 1]
+        rewards_np = np.array([self.memory.rewards[i] for i in idxs])
+        r_min = rewards_np.min()
+        r_max = rewards_np.max()
+        # Avoid division by zero
+        r_scaled = 2 * (rewards_np - r_min) / (r_max - r_min + 1e-8) - 1
+        normalized_rewards = np.tanh(r_scaled)
+        rewards = T.as_tensor(normalized_rewards, dtype=T.float32, device=self.device).unsqueeze(-1)
+        """
+        
+        # Critic update, Soft Q-Learning Objective: to ensure high-entropy actions for exploration 
+        with T.no_grad():
+            next_logits = self.policy(next_states)
+            next_dist = Categorical(logits=next_logits)
+            next_probs = next_dist.probs
+            next_log_probs = next_dist.logits - T.logsumexp(next_dist.logits, dim=-1, keepdim=True)
+            q1_next = self.q1_target(next_states)
+            q2_next = self.q2_target(next_states)
+            # Soft Policy Evaluation 
+            min_q_next = T.min(q1_next, q2_next)
+            next_value = (next_probs * (min_q_next - self.alpha * next_log_probs)).sum(dim=-1, keepdim=True)
+            target = rewards + self.gamma * (1 - dones) * next_value
+
+        q1 = self.q1(states).gather(1, actions)
+        q2 = self.q2(states).gather(1, actions)
+
+        # Losses of both Q-functions
+        q1_loss = nn.MSELoss()(q1, target)
+        q2_loss = nn.MSELoss()(q2, target)
+
+        self.q1_opt.zero_grad()
+        q1_loss.backward()
+        self.q1_opt.step()
+        self.q2_opt.zero_grad()
+        q2_loss.backward()
+        self.q2_opt.step()
+
+        # Policy/Actor Objective
+        logits = self.policy(states)
+        dist = Categorical(logits=logits)
+        probs = dist.probs
+        log_probs = dist.logits - T.logsumexp(dist.logits, dim=-1, keepdim=True)
+        q1_policy = self.q1(states)
+        q2_policy = self.q2(states)
+        min_q_policy = T.min(q1_policy, q2_policy)
+        # Slightly different policy loss for discrete actions
+        policy_loss = (probs * (self.alpha * log_probs - min_q_policy)).sum(dim=-1).mean()
+
+        # Temperature to update Alpha
+        alpha_loss = -(self.log_alpha * (log_probs + self.target_entropy).detach()).mean()
+        self.alpha_opt.zero_grad()
+        alpha_loss.backward()
+        T.nn.utils.clip_grad_norm_([self.log_alpha], max_norm=1.0)  # Gradient clipping
+        self.alpha_opt.step()
+        self.alpha = self.log_alpha.exp().item()  
+
+        self.policy_opt.zero_grad()
+        policy_loss.backward()
+        T.nn.utils.clip_grad_norm_([self.log_alpha], max_norm=1.0)  # Gradient clipping
+        self.policy_opt.step()
+
+        # Target network update
+        for target_param, param in zip(self.q1_target.parameters(), self.q1.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+        for target_param, param in zip(self.q2_target.parameters(), self.q2.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+
+        return policy_loss.item()
+
+# Actor/Policy network 
+# Typical SAC Actor network is used to output a Gaussian distribution of a state
+# Here, we adapt it for discrete actions using a Categorical distribution, as the ATARI environment is discrete
+# The policy outputs logits for each discrete action.
+
+# From: https://ch.mathworks.com/help/reinforcement-learning/ug/soft-actor-critic-agents.html
+# The actor takes the current observation and generates a categorical distribution, in which each possible action is associated with a probability. 
+
+class CategoricalActor(nn.Module):
+    def __init__(self, obs_shape: tuple, action_dim: int, hidden: int):
+        super().__init__()
+        c, h, w = obs_shape
+        self.cnn = nn.Sequential(
+            nn.Conv2d(c, 16, kernel_size=8, stride=4),
+            nn.ReLU(),
+            nn.Conv2d(16, 32, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Flatten()
+        )
+        with T.no_grad():
+            cnn_output_dim = self.cnn(T.zeros(1, c, h, w)).shape[1]
+        self.fc = nn.Sequential(
+            nn.Linear(cnn_output_dim, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, action_dim)
+        )
+
+    def forward(self, state: T.Tensor):
+        if state.dim() == 3:
+            state = state.unsqueeze(0)
+        cnn_out = self.cnn(state)
+        logits = self.fc(cnn_out)
+        return logits
+
+# Q-network for discrete actions 
+class QNetwork(nn.Module):
+    def __init__(self, obs_shape: tuple, action_dim: int, hidden: int):
+        super().__init__()
+        c, h, w = obs_shape
+        self.cnn = nn.Sequential(
+            nn.Conv2d(c, 16, kernel_size=8, stride=4),
+            nn.ReLU(),
+            nn.Conv2d(16, 32, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Flatten()
+        )
+        with T.no_grad():
+            cnn_output_dim = self.cnn(T.zeros(1, c, h, w)).shape[1]
+        self.net = nn.Sequential(
+            nn.Linear(cnn_output_dim, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, action_dim)
+        )
+
+    def forward(self, state: T.Tensor):
+        if state.dim() == 3:
+            state = state.unsqueeze(0)
+        cnn_out = self.cnn(state)
+        return self.net(cnn_out)
+
+class Memory:
+    def __init__(self):
+        self.states, self.actions, self.rewards, self.dones, self.next_states = [], [], [], [], []
+    def store(self, s, a, r, d, ns):
+        self.states.append(np.asarray(s, dtype=np.float32))
+        self.actions.append(np.asarray(a, dtype=np.float32))
+        self.rewards.append(float(r))
+        self.dones.append(float(d))
+        self.next_states.append(np.asarray(ns, dtype=np.float32))
+    def clear(self):
+        self.states, self.actions, self.rewards, self.dones, self.next_states = [], [], [], [], []
+
+
+
+
+
+class ObservationNorm:
+
+
+    def normalize(self, x):
+        return (x - x.mean()) / (x.std(unbiased=False) + 1e-8)  # We add epsilon to make sure that we don't
+        # divide through zero.

Observation_norm_SAC/Obser_norm_in_batch/sac_model_cnn_in_batch.py

@@ -0,0 +1,206 @@
+import ale_py
+import gymnasium as gym
+import sys
+import numpy as np
+from sac_helpers_cnn_in_batch import *
+from gymnasium.spaces import Box
+import cv2
+import matplotlib.pyplot as plt
+
+def preprocess(obs):
+    obs = cv2.cvtColor(obs, cv2.COLOR_RGB2GRAY)
+    obs = cv2.resize(obs, (84, 84), interpolation=cv2.INTER_AREA)
+    obs = np.expand_dims(obs, axis=0) 
+    return obs.astype(np.float32) / 255.0
+"""
+def main() -> int:
+    episode = 0
+    total_return = 0
+    ep_return = 0
+    steps = 100
+    batches = 100
+    avg_returns = []
+    avg_losses = []
+
+    env = gym.make("ALE/Pacman-v5")
+    # Initialize CNN with a dummy observation (to get correct input shape)
+    obs, _ = env.reset()
+    dummy_obs_space = Box(low=0.0, high=1.0, shape=preprocess(obs).shape)
+
+    agent = Agent(
+        obs_space=dummy_obs_space,
+        action_space=env.action_space,
+        hidden=64,
+        gamma=0.99,
+        lr=3e-4,
+        alpha=0.2,
+        seed=70,
+        batch_size=32,
+        tau=0.005
+    )
+
+    try:
+        obs, info = env.reset(seed=42)
+        state = preprocess(obs)
+
+        for update in range(1, batches + 1):
+            batch_loss = []
+
+            for t in range(steps):
+                action = agent.choose_action(state)
+                next_obs, reward, terminated, truncated, info = env.step(action)
+                done = terminated or truncated
+                next_state = preprocess(next_obs)
+
+                agent.remember(state, action, reward, done, next_state)
+
+                ep_return += reward
+                state = next_state
+
+                if done:
+                    episode += 1
+                    total_return += ep_return
+                    print(f"Episode {episode} return: {ep_return:.2f}")
+                    ep_return = 0
+                    obs, info = env.reset()
+                    state = preprocess(obs)
+
+            loss = agent.vanilla_sac_update()
+            batch_loss.append(loss) 
+
+            avg_ret = (total_return / episode) if episode else 0
+            avg_returns.append(avg_ret)
+            avg_losses.append(np.mean(batch_loss)) 
+
+            print(f"Update {update}: episodes={episode}, avg_return={avg_ret:.2f}, avg_loss={np.mean(batch_loss):.4f}")
+            
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        return 1
+    finally:
+        avg = total_return / episode if episode else 0
+        print(f"\nEpisodes: {episode}, Avg return: {avg:.3f}")
+
+        # Plot learning curve
+        plt.figure(figsize=(10, 4))
+        plt.subplot(1, 2, 1)
+        plt.plot(avg_returns)
+        plt.xlabel("Update")
+        plt.ylabel("Average Return")
+        plt.title("SAC Learning Curve")
+        plt.grid()
+
+        plt.subplot(1, 2, 2)
+        plt.plot(avg_losses)
+        plt.xlabel("Update")
+        plt.ylabel("Average Loss")
+        plt.title("Average Loss Curve")
+        plt.grid()
+
+        plt.tight_layout()
+        plt.show()
+        env.close()
+
+    return 0
+"""
+
+def run_training(seed=42):
+    episode = 0
+    total_return = 0
+    ep_return = 0
+    steps = 100
+    batches = 100
+    avg_returns = []
+    avg_losses = []
+
+    env = gym.make("ALE/Pacman-v5")
+    obs, _ = env.reset()
+    dummy_obs_space = Box(low=0.0, high=1.0, shape=preprocess(obs).shape)
+
+    agent = Agent(
+        obs_space=dummy_obs_space,
+        action_space=env.action_space,
+        hidden=64,
+        gamma=0.99,
+        lr=3e-4,
+        alpha=0.2,
+        seed=seed,
+        batch_size=32,
+        tau=0.005
+    )
+
+    obs, info = env.reset(seed=seed)
+    state = preprocess(obs)
+
+    for update in range(1, batches + 1):
+        batch_loss = []
+        for t in range(steps):
+            action = agent.choose_action(state)
+            next_obs, reward, terminated, truncated, info = env.step(action)
+            done = terminated or truncated
+            next_state = preprocess(next_obs)
+
+            agent.remember(state, action, reward, done, next_state)
+
+            ep_return += reward
+            state = next_state
+
+            if done:
+                episode += 1
+                total_return += ep_return
+                ep_return = 0
+                obs, info = env.reset()
+                state = preprocess(obs)
+
+        loss = agent.vanilla_sac_update()
+        #loss = agent.update_reward_gradient_clipping()
+        batch_loss.append(loss)
+
+        avg_ret = (total_return / episode) if episode else 0
+        avg_returns.append(avg_ret)
+        avg_losses.append(np.mean(batch_loss))
+
+    env.close()
+    return avg_returns, avg_losses
+
+def main() -> int:
+    num_runs = 5
+    all_returns = []
+    all_losses = []
+
+    for run in range(num_runs):
+        print(f"Starting run {run+1}/{num_runs}")
+        avg_returns, avg_losses = run_training(seed=42 + run)
+        all_returns.append(avg_returns)
+        all_losses.append(avg_losses)
+
+    # Convert to numpy arrays for easy averaging
+    all_returns = np.array(all_returns)
+    all_losses = np.array(all_losses)
+
+    mean_returns = np.mean(all_returns, axis=0)
+    mean_losses = np.mean(all_losses, axis=0)
+
+    # Plot averaged learning curves
+    plt.figure(figsize=(10, 4))
+    plt.subplot(1, 2, 1)
+    plt.plot(mean_returns)
+    plt.xlabel("Update")
+    plt.ylabel("Average Return")
+    plt.title(f"SAC Learning Curve (avg over {num_runs} runs)")
+    plt.grid()
+
+    plt.subplot(1, 2, 2)
+    plt.plot(mean_losses)
+    plt.xlabel("Update")
+    plt.ylabel("Average Loss")
+    plt.title(f"Average Loss Curve (avg over {num_runs} runs)")
+    plt.grid()
+
+    plt.tight_layout()
+    plt.savefig("Observation norm." + "_in_batch_.png")
+
+    return 0
+
+if __name__ == "__main__":
+    raise SystemExit(main())

Observation_norm_SAC/Obser_norm_running_average/sac_helpers_cnn_running_average.py

@@ -0,0 +1,350 @@
+import numpy as np
+import torch as T
+import torch.nn as nn
+import torch.optim as optim
+from torch.distributions import Categorical
+
+class Agent:
+    def __init__(self, obs_space, action_space, hidden, gamma, lr, alpha, seed, batch_size, tau=0.005):
+        if seed is not None:
+            np.random.seed(seed)
+            T.manual_seed(seed)
+        
+        # Use GPU if available
+
+        self.device = T.device('cuda:0' if T.cuda.is_available() else 'cpu')
+        self.action_dim = int(getattr(action_space, "n", action_space.n))  # Use .n for Discrete
+        self.obs_shape = obs_space.shape
+
+        self.gamma, self.tau, self.batch_size = gamma, tau, batch_size
+        # Make alpha learnable (adjust entropy based on reward magnitude)
+        self.target_entropy = -float(self.action_dim)
+        self.log_alpha = T.tensor(np.log(alpha), requires_grad=True, device=self.device)
+        self.alpha = np.exp(self.log_alpha.item()) 
+        self.alpha_opt = optim.Adam([self.log_alpha], lr=lr)
+        self.observeNorm = ObservationNorm(self.obs_shape)
+
+
+        self.policy = CategoricalActor(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q1 = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q2 = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q1_target = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q2_target = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q1_target.load_state_dict(self.q1.state_dict())
+        self.q2_target.load_state_dict(self.q2.state_dict())
+
+        self.policy_opt = optim.Adam(self.policy.parameters(), lr=lr)
+        self.q1_opt = optim.Adam(self.q1.parameters(), lr=lr)
+        self.q2_opt = optim.Adam(self.q2.parameters(), lr=lr)
+        self.memory = Memory()
+
+    def choose_action(self, observation, eval_mode=False):
+        state = T.as_tensor(observation, dtype=T.float32, device=self.device)
+
+        state = self.observeNorm.normalize(state.unsqueeze(0))
+        state = state.squeeze(0)
+        with T.no_grad():
+            logits = self.policy(state.unsqueeze(0))
+            dist = Categorical(logits=logits)
+            if eval_mode:
+                action = logits.argmax(dim=-1)
+            else:
+                action = dist.sample()
+        return int(action.item())
+
+    def remember(self, state, action, reward, done, next_state):
+
+        state = T.as_tensor(state, dtype=T.float32, device=self.device)
+        if state.dim() == 3:  # [C,H,W]
+            state = state.unsqueeze(0)  # [1,C,H,W]
+        self.observeNorm.update(state)
+        state = self.observeNorm.normalize(state)
+        state = state.squeeze(0)
+        #next_state also needs normalization
+
+        next_state = T.as_tensor(next_state, dtype=T.float32, device=self.device)
+        if next_state.dim() == 3:  # [C,H,W]
+            next_state = next_state.unsqueeze(0)  # [1,C,H,W]
+        self.observeNorm.update(next_state)
+        next_state = self.observeNorm.normalize(next_state)
+        next_state = next_state.squeeze(0)
+
+
+        self.memory.store(state, action, reward, done, next_state)
+
+    def vanilla_sac_update(self):
+        if len(self.memory.states) < self.batch_size:
+            return 0.0
+        
+        # Mini-batch sampling
+        idxs = np.random.choice(len(self.memory.states), self.batch_size, replace=False)
+        states = T.as_tensor(np.array([self.memory.states[i] for i in idxs]), dtype=T.float32, device=self.device)
+        actions = T.as_tensor(np.array([self.memory.actions[i] for i in idxs]), dtype=T.int64, device=self.device).unsqueeze(-1)
+        rewards = T.as_tensor(np.array([self.memory.rewards[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        dones = T.as_tensor(np.array([self.memory.dones[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        next_states = T.as_tensor(np.array([self.memory.next_states[i] for i in idxs]), dtype=T.float32, device=self.device)
+
+        # Critic update, Soft Q-Learning Objective: to ensure high-entropy actions for exploration 
+        with T.no_grad():
+            next_logits = self.policy(next_states)
+            next_dist = Categorical(logits=next_logits)
+            next_probs = next_dist.probs
+            next_log_probs = next_dist.logits - T.logsumexp(next_dist.logits, dim=-1, keepdim=True)
+            q1_next = self.q1_target(next_states)
+            q2_next = self.q2_target(next_states)
+            # Soft Policy Evaluation 
+            min_q_next = T.min(q1_next, q2_next)
+            next_value = (next_probs * (min_q_next - self.alpha * next_log_probs)).sum(dim=-1, keepdim=True)
+            target = rewards + self.gamma * (1 - dones) * next_value
+
+        q1 = self.q1(states).gather(1, actions)
+        q2 = self.q2(states).gather(1, actions)
+
+        # Losses of both Q-functions
+        q1_loss = nn.MSELoss()(q1, target)
+        q2_loss = nn.MSELoss()(q2, target)
+
+        self.q1_opt.zero_grad()
+        q1_loss.backward()
+        self.q1_opt.step()
+        self.q2_opt.zero_grad()
+        q2_loss.backward()
+        self.q2_opt.step()
+
+        # Policy/Actor Objective
+        logits = self.policy(states)
+        dist = Categorical(logits=logits)
+        probs = dist.probs
+        log_probs = dist.logits - T.logsumexp(dist.logits, dim=-1, keepdim=True)
+        q1_policy = self.q1(states)
+        q2_policy = self.q2(states)
+        min_q_policy = T.min(q1_policy, q2_policy)
+        # Slightly different policy loss for discrete actions
+        policy_loss = (probs * (self.alpha * log_probs - min_q_policy)).sum(dim=-1).mean()
+
+        # Temperature to update Alpha
+        alpha_loss = -(self.log_alpha * (log_probs + self.target_entropy).detach()).mean()
+        self.alpha_opt.zero_grad()
+        alpha_loss.backward()
+        self.alpha_opt.step()
+        self.alpha = self.log_alpha.exp().item()  
+
+        self.policy_opt.zero_grad()
+        policy_loss.backward()
+        self.policy_opt.step()
+
+        # Target network update
+        for target_param, param in zip(self.q1_target.parameters(), self.q1.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+        for target_param, param in zip(self.q2_target.parameters(), self.q2.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+
+        return policy_loss.item()
+    
+    def update_reward_gradient_clipping(self):
+        if len(self.memory.states) < self.batch_size:
+            return 0.0
+        
+        # Mini-batch sampling
+        idxs = np.random.choice(len(self.memory.states), self.batch_size, replace=False)
+        states = T.as_tensor(np.array([self.memory.states[i] for i in idxs]), dtype=T.float32, device=self.device)
+        actions = T.as_tensor(np.array([self.memory.actions[i] for i in idxs]), dtype=T.int64, device=self.device).unsqueeze(-1)
+        rewards = T.as_tensor(np.array([self.memory.rewards[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        dones = T.as_tensor(np.array([self.memory.dones[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        next_states = T.as_tensor(np.array([self.memory.next_states[i] for i in idxs]), dtype=T.float32, device=self.device)
+        
+        """
+        # Min-max normalization and tanh scaling to [-1, 1]
+        rewards_np = np.array([self.memory.rewards[i] for i in idxs])
+        r_min = rewards_np.min()
+        r_max = rewards_np.max()
+        # Avoid division by zero
+        r_scaled = 2 * (rewards_np - r_min) / (r_max - r_min + 1e-8) - 1
+        normalized_rewards = np.tanh(r_scaled)
+        rewards = T.as_tensor(normalized_rewards, dtype=T.float32, device=self.device).unsqueeze(-1)
+        """
+        
+        # Critic update, Soft Q-Learning Objective: to ensure high-entropy actions for exploration 
+        with T.no_grad():
+            next_logits = self.policy(next_states)
+            next_dist = Categorical(logits=next_logits)
+            next_probs = next_dist.probs
+            next_log_probs = next_dist.logits - T.logsumexp(next_dist.logits, dim=-1, keepdim=True)
+            q1_next = self.q1_target(next_states)
+            q2_next = self.q2_target(next_states)
+            # Soft Policy Evaluation 
+            min_q_next = T.min(q1_next, q2_next)
+            next_value = (next_probs * (min_q_next - self.alpha * next_log_probs)).sum(dim=-1, keepdim=True)
+            target = rewards + self.gamma * (1 - dones) * next_value
+
+        q1 = self.q1(states).gather(1, actions)
+        q2 = self.q2(states).gather(1, actions)
+
+        # Losses of both Q-functions
+        q1_loss = nn.MSELoss()(q1, target)
+        q2_loss = nn.MSELoss()(q2, target)
+
+        self.q1_opt.zero_grad()
+        q1_loss.backward()
+        self.q1_opt.step()
+        self.q2_opt.zero_grad()
+        q2_loss.backward()
+        self.q2_opt.step()
+
+        # Policy/Actor Objective
+        logits = self.policy(states)
+        dist = Categorical(logits=logits)
+        probs = dist.probs
+        log_probs = dist.logits - T.logsumexp(dist.logits, dim=-1, keepdim=True)
+        q1_policy = self.q1(states)
+        q2_policy = self.q2(states)
+        min_q_policy = T.min(q1_policy, q2_policy)
+        # Slightly different policy loss for discrete actions
+        policy_loss = (probs * (self.alpha * log_probs - min_q_policy)).sum(dim=-1).mean()
+
+        # Temperature to update Alpha
+        alpha_loss = -(self.log_alpha * (log_probs + self.target_entropy).detach()).mean()
+        self.alpha_opt.zero_grad()
+        alpha_loss.backward()
+        T.nn.utils.clip_grad_norm_([self.log_alpha], max_norm=1.0)  # Gradient clipping
+        self.alpha_opt.step()
+        self.alpha = self.log_alpha.exp().item()  
+
+        self.policy_opt.zero_grad()
+        policy_loss.backward()
+        T.nn.utils.clip_grad_norm_([self.log_alpha], max_norm=1.0)  # Gradient clipping
+        self.policy_opt.step()
+
+        # Target network update
+        for target_param, param in zip(self.q1_target.parameters(), self.q1.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+        for target_param, param in zip(self.q2_target.parameters(), self.q2.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+
+        return policy_loss.item()
+
+
+
+# Actor/Policy network 
+# Typical SAC Actor network is used to output a Gaussian distribution of a state
+# Here, we adapt it for discrete actions using a Categorical distribution, as the ATARI environment is discrete
+# The policy outputs logits for each discrete action.
+
+# From: https://ch.mathworks.com/help/reinforcement-learning/ug/soft-actor-critic-agents.html
+# The actor takes the current observation and generates a categorical distribution, in which each possible action is associated with a probability. 
+
+class CategoricalActor(nn.Module):
+    def __init__(self, obs_shape: tuple, action_dim: int, hidden: int):
+        super().__init__()
+        c, h, w = obs_shape
+        self.cnn = nn.Sequential(
+            nn.Conv2d(c, 16, kernel_size=8, stride=4),
+            nn.ReLU(),
+            nn.Conv2d(16, 32, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Flatten()
+        )
+        with T.no_grad():
+            cnn_output_dim = self.cnn(T.zeros(1, c, h, w)).shape[1]
+        self.fc = nn.Sequential(
+            nn.Linear(cnn_output_dim, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, action_dim)
+        )
+
+    def forward(self, state: T.Tensor):
+        if state.dim() == 3:
+            state = state.unsqueeze(0)
+        cnn_out = self.cnn(state)
+        logits = self.fc(cnn_out)
+        return logits
+
+# Q-network for discrete actions 
+class QNetwork(nn.Module):
+    def __init__(self, obs_shape: tuple, action_dim: int, hidden: int):
+        super().__init__()
+        c, h, w = obs_shape
+        self.cnn = nn.Sequential(
+            nn.Conv2d(c, 16, kernel_size=8, stride=4),
+            nn.ReLU(),
+            nn.Conv2d(16, 32, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Flatten()
+        )
+        with T.no_grad():
+            cnn_output_dim = self.cnn(T.zeros(1, c, h, w)).shape[1]
+        self.net = nn.Sequential(
+            nn.Linear(cnn_output_dim, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, action_dim)
+        )
+
+    def forward(self, state: T.Tensor):
+        if state.dim() == 3:
+            state = state.unsqueeze(0)
+        cnn_out = self.cnn(state)
+        return self.net(cnn_out)
+
+class Memory:
+    def __init__(self):
+        self.states, self.actions, self.rewards, self.dones, self.next_states = [], [], [], [], []
+    def store(self, s, a, r, d, ns):
+        self.states.append(np.asarray(s, dtype=np.float32))
+        self.actions.append(np.asarray(a, dtype=np.float32))
+        self.rewards.append(float(r))
+        self.dones.append(float(d))
+        self.next_states.append(np.asarray(ns, dtype=np.float32))
+    def clear(self):
+        self.states, self.actions, self.rewards, self.dones, self.next_states = [], [], [], [], []
+
+
+
+
+
+
+class ObservationNorm:
+    def __init__(self, shape):
+        c, h, w = shape
+        self.mean = T.zeros((c, 1, 1))
+        self.var = T.ones((c, 1, 1))
+        self.count = 1e-4
+
+    def update(self, x: T.Tensor):
+        batch_mean = x.mean(dim=[0, 2, 3], keepdim=True)
+        batch_var = x.var(dim=[0, 2, 3], keepdim=True, unbiased=False)
+        batch_count = x.shape[0]
+        self._update_from_moments(batch_mean, batch_var, batch_count)
+
+    def _update_from_moments(self, batch_mean, batch_var, batch_count):
+        delta = batch_mean - self.mean
+        total_count = self.count + batch_count
+
+        # Update mean
+        new_mean = self.mean + delta * batch_count / total_count
+
+        # Update variance
+        m_a = self.var * self.count  # scaled old variance
+        m_b = batch_var * batch_count  # scaled batch variance
+        M2 = m_a + m_b + (delta ** 2) * (self.count * batch_count / total_count)
+        new_var = M2 / total_count
+
+        # Assign updates
+        self.mean = new_mean
+        self.var = new_var
+        self.count = total_count
+
+    def normalize(self, x:T.Tensor):
+        return (x - self.mean) / (T.sqrt(self.var) + 1e-8) # We add epsilon to make sure that we don't
+                                                            # divide through zero.
+
+
+
+
+
+
+
+
+
+
+
+

Observation_norm_SAC/Obser_norm_running_average/sac_model_cnn_running_average.py

@@ -0,0 +1,207 @@
+import ale_py
+import gymnasium as gym
+import sys
+import numpy as np
+from sac_helpers_cnn_running_average import *
+from gymnasium.spaces import Box
+import cv2
+import matplotlib.pyplot as plt
+
+def preprocess(obs):
+    obs = cv2.cvtColor(obs, cv2.COLOR_RGB2GRAY)
+    obs = cv2.resize(obs, (84, 84), interpolation=cv2.INTER_AREA)
+    obs = np.expand_dims(obs, axis=0) 
+    return obs.astype(np.float32) / 255.0
+"""
+def main() -> int:
+    episode = 0
+    total_return = 0
+    ep_return = 0
+    steps = 100
+    batches = 100
+    avg_returns = []
+    avg_losses = []
+
+    env = gym.make("ALE/Pacman-v5")
+    # Initialize CNN with a dummy observation (to get correct input shape)
+    obs, _ = env.reset()
+    dummy_obs_space = Box(low=0.0, high=1.0, shape=preprocess(obs).shape)
+
+    agent = Agent(
+        obs_space=dummy_obs_space,
+        action_space=env.action_space,
+        hidden=64,
+        gamma=0.99,
+        lr=3e-4,
+        alpha=0.2,
+        seed=70,
+        batch_size=32,
+        tau=0.005
+    )
+
+    try:
+        obs, info = env.reset(seed=42)
+        state = preprocess(obs)
+
+        for update in range(1, batches + 1):
+            batch_loss = []
+
+            for t in range(steps):
+                action = agent.choose_action(state)
+                next_obs, reward, terminated, truncated, info = env.step(action)
+                done = terminated or truncated
+                next_state = preprocess(next_obs)
+
+                agent.remember(state, action, reward, done, next_state)
+
+                ep_return += reward
+                state = next_state
+
+                if done:
+                    episode += 1
+                    total_return += ep_return
+                    print(f"Episode {episode} return: {ep_return:.2f}")
+                    ep_return = 0
+                    obs, info = env.reset()
+                    state = preprocess(obs)
+
+            loss = agent.vanilla_sac_update()
+            batch_loss.append(loss) 
+
+            avg_ret = (total_return / episode) if episode else 0
+            avg_returns.append(avg_ret)
+            avg_losses.append(np.mean(batch_loss)) 
+
+            print(f"Update {update}: episodes={episode}, avg_return={avg_ret:.2f}, avg_loss={np.mean(batch_loss):.4f}")
+            
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        return 1
+    finally:
+        avg = total_return / episode if episode else 0
+        print(f"\nEpisodes: {episode}, Avg return: {avg:.3f}")
+
+        # Plot learning curve
+        plt.figure(figsize=(10, 4))
+        plt.subplot(1, 2, 1)
+        plt.plot(avg_returns)
+        plt.xlabel("Update")
+        plt.ylabel("Average Return")
+        plt.title("SAC Learning Curve")
+        plt.grid()
+
+        plt.subplot(1, 2, 2)
+        plt.plot(avg_losses)
+        plt.xlabel("Update")
+        plt.ylabel("Average Loss")
+        plt.title("Average Loss Curve")
+        plt.grid()
+
+        plt.tight_layout()
+        plt.show()
+        env.close()
+
+    return 0
+"""
+
+def run_training(seed=42):
+    episode = 0
+    total_return = 0
+    ep_return = 0
+    steps = 100
+    batches = 100
+    avg_returns = []
+    avg_losses = []
+
+    env = gym.make("ALE/Pacman-v5")
+    obs, _ = env.reset()
+    dummy_obs_space = Box(low=0.0, high=1.0, shape=preprocess(obs).shape)
+
+    agent = Agent(
+        obs_space=dummy_obs_space,
+        action_space=env.action_space,
+        hidden=64,
+        gamma=0.99,
+        lr=3e-4,
+        alpha=0.2,
+        seed=seed,
+        batch_size=32,
+        tau=0.005
+    )
+
+    obs, info = env.reset(seed=seed)
+    state = preprocess(obs)
+
+    for update in range(1, batches + 1):
+        batch_loss = []
+        for t in range(steps):
+            action = agent.choose_action(state)
+            next_obs, reward, terminated, truncated, info = env.step(action)
+            done = terminated or truncated
+            next_state = preprocess(next_obs)
+
+            agent.remember(state, action, reward, done, next_state)
+
+            ep_return += reward
+            state = next_state
+
+            if done:
+                episode += 1
+                total_return += ep_return
+                ep_return = 0
+                obs, info = env.reset()
+                state = preprocess(obs)
+
+        loss = agent.vanilla_sac_update()
+        #loss = agent.update_reward_gradient_clipping()
+        batch_loss.append(loss)
+
+        avg_ret = (total_return / episode) if episode else 0
+        avg_returns.append(avg_ret)
+        avg_losses.append(np.mean(batch_loss))
+
+    env.close()
+    return avg_returns, avg_losses
+
+def main() -> int:
+    num_runs = 5
+    all_returns = []
+    all_losses = []
+
+    for run in range(num_runs):
+        print(f"Starting run {run+1}/{num_runs}")
+        avg_returns, avg_losses = run_training(seed=42 + run)
+        all_returns.append(avg_returns)
+        all_losses.append(avg_losses)
+
+    # Convert to numpy arrays for easy averaging
+    all_returns = np.array(all_returns)
+    all_losses = np.array(all_losses)
+
+    mean_returns = np.mean(all_returns, axis=0)
+    mean_losses = np.mean(all_losses, axis=0)
+
+    # Plot averaged learning curves
+    plt.figure(figsize=(10, 4))
+    plt.subplot(1, 2, 1)
+    plt.plot(mean_returns)
+    plt.xlabel("Update")
+    plt.ylabel("Average Return")
+    plt.title(f"SAC Learning Curve (avg over {num_runs} runs)")
+    plt.grid()
+
+    plt.subplot(1, 2, 2)
+    plt.plot(mean_losses)
+    plt.xlabel("Update")
+    plt.ylabel("Average Loss")
+    plt.title(f"Average Loss Curve (avg over {num_runs} runs)")
+    plt.grid()
+
+    plt.tight_layout()
+    plt.savefig("Observation norm." + "_running_average_.png")
+
+
+    return 0
+
+if __name__ == "__main__":
+    raise SystemExit(main())

Observation_norm_SAC/sac_helpers_cnn.py

@@ -0,0 +1,274 @@
+import numpy as np
+import torch as T
+import torch.nn as nn
+import torch.optim as optim
+from torch.distributions import Categorical
+
+class Agent:
+    def __init__(self, obs_space, action_space, hidden, gamma, lr, alpha, seed, batch_size, tau=0.005):
+        if seed is not None:
+            np.random.seed(seed)
+            T.manual_seed(seed)
+        
+        # Use GPU if available
+
+        self.device = T.device('cuda:0' if T.cuda.is_available() else 'cpu')
+        self.action_dim = int(getattr(action_space, "n", action_space.n))  # Use .n for Discrete
+        self.obs_shape = obs_space.shape
+
+        self.gamma, self.tau, self.batch_size = gamma, tau, batch_size
+        # Make alpha learnable (adjust entropy based on reward magnitude)
+        self.target_entropy = -float(self.action_dim)
+        self.log_alpha = T.tensor(np.log(alpha), requires_grad=True, device=self.device)
+        self.alpha = np.exp(self.log_alpha.item()) 
+        self.alpha_opt = optim.Adam([self.log_alpha], lr=lr)
+
+        self.policy = CategoricalActor(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q1 = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q2 = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q1_target = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q2_target = QNetwork(self.obs_shape, self.action_dim, hidden).to(self.device)
+        self.q1_target.load_state_dict(self.q1.state_dict())
+        self.q2_target.load_state_dict(self.q2.state_dict())
+
+        self.policy_opt = optim.Adam(self.policy.parameters(), lr=lr)
+        self.q1_opt = optim.Adam(self.q1.parameters(), lr=lr)
+        self.q2_opt = optim.Adam(self.q2.parameters(), lr=lr)
+        self.memory = Memory()
+
+    def choose_action(self, observation, eval_mode=False):
+        state = T.as_tensor(observation, dtype=T.float32, device=self.device)
+        with T.no_grad():
+            logits = self.policy(state.unsqueeze(0))
+            dist = Categorical(logits=logits)
+            if eval_mode:
+                action = logits.argmax(dim=-1)
+            else:
+                action = dist.sample()
+        return int(action.item())
+
+    def remember(self, state, action, reward, done, next_state):
+        self.memory.store(state, action, reward, done, next_state)
+
+    def vanilla_sac_update(self):
+        if len(self.memory.states) < self.batch_size:
+            return 0.0
+        
+        # Mini-batch sampling
+        idxs = np.random.choice(len(self.memory.states), self.batch_size, replace=False)
+        states = T.as_tensor(np.array([self.memory.states[i] for i in idxs]), dtype=T.float32, device=self.device)
+        actions = T.as_tensor(np.array([self.memory.actions[i] for i in idxs]), dtype=T.int64, device=self.device).unsqueeze(-1)
+        rewards = T.as_tensor(np.array([self.memory.rewards[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        dones = T.as_tensor(np.array([self.memory.dones[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        next_states = T.as_tensor(np.array([self.memory.next_states[i] for i in idxs]), dtype=T.float32, device=self.device)
+
+        # Critic update, Soft Q-Learning Objective: to ensure high-entropy actions for exploration 
+        with T.no_grad():
+            next_logits = self.policy(next_states)
+            next_dist = Categorical(logits=next_logits)
+            next_probs = next_dist.probs
+            next_log_probs = next_dist.logits - T.logsumexp(next_dist.logits, dim=-1, keepdim=True)
+            q1_next = self.q1_target(next_states)
+            q2_next = self.q2_target(next_states)
+            # Soft Policy Evaluation 
+            min_q_next = T.min(q1_next, q2_next)
+            next_value = (next_probs * (min_q_next - self.alpha * next_log_probs)).sum(dim=-1, keepdim=True)
+            target = rewards + self.gamma * (1 - dones) * next_value
+
+        q1 = self.q1(states).gather(1, actions)
+        q2 = self.q2(states).gather(1, actions)
+
+        # Losses of both Q-functions
+        q1_loss = nn.MSELoss()(q1, target)
+        q2_loss = nn.MSELoss()(q2, target)
+
+        self.q1_opt.zero_grad()
+        q1_loss.backward()
+        self.q1_opt.step()
+        self.q2_opt.zero_grad()
+        q2_loss.backward()
+        self.q2_opt.step()
+
+        # Policy/Actor Objective
+        logits = self.policy(states)
+        dist = Categorical(logits=logits)
+        probs = dist.probs
+        log_probs = dist.logits - T.logsumexp(dist.logits, dim=-1, keepdim=True)
+        q1_policy = self.q1(states)
+        q2_policy = self.q2(states)
+        min_q_policy = T.min(q1_policy, q2_policy)
+        # Slightly different policy loss for discrete actions
+        policy_loss = (probs * (self.alpha * log_probs - min_q_policy)).sum(dim=-1).mean()
+
+        # Temperature to update Alpha
+        alpha_loss = -(self.log_alpha * (log_probs + self.target_entropy).detach()).mean()
+        self.alpha_opt.zero_grad()
+        alpha_loss.backward()
+        self.alpha_opt.step()
+        self.alpha = self.log_alpha.exp().item()  
+
+        self.policy_opt.zero_grad()
+        policy_loss.backward()
+        self.policy_opt.step()
+
+        # Target network update
+        for target_param, param in zip(self.q1_target.parameters(), self.q1.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+        for target_param, param in zip(self.q2_target.parameters(), self.q2.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+
+        return policy_loss.item()
+    
+    def update_reward_gradient_clipping(self):
+        if len(self.memory.states) < self.batch_size:
+            return 0.0
+        
+        # Mini-batch sampling
+        idxs = np.random.choice(len(self.memory.states), self.batch_size, replace=False)
+        states = T.as_tensor(np.array([self.memory.states[i] for i in idxs]), dtype=T.float32, device=self.device)
+        actions = T.as_tensor(np.array([self.memory.actions[i] for i in idxs]), dtype=T.int64, device=self.device).unsqueeze(-1)
+        rewards = T.as_tensor(np.array([self.memory.rewards[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        dones = T.as_tensor(np.array([self.memory.dones[i] for i in idxs]), dtype=T.float32, device=self.device).unsqueeze(-1)
+        next_states = T.as_tensor(np.array([self.memory.next_states[i] for i in idxs]), dtype=T.float32, device=self.device)
+        
+        """
+        # Min-max normalization and tanh scaling to [-1, 1]
+        rewards_np = np.array([self.memory.rewards[i] for i in idxs])
+        r_min = rewards_np.min()
+        r_max = rewards_np.max()
+        # Avoid division by zero
+        r_scaled = 2 * (rewards_np - r_min) / (r_max - r_min + 1e-8) - 1
+        normalized_rewards = np.tanh(r_scaled)
+        rewards = T.as_tensor(normalized_rewards, dtype=T.float32, device=self.device).unsqueeze(-1)
+        """
+        
+        # Critic update, Soft Q-Learning Objective: to ensure high-entropy actions for exploration 
+        with T.no_grad():
+            next_logits = self.policy(next_states)
+            next_dist = Categorical(logits=next_logits)
+            next_probs = next_dist.probs
+            next_log_probs = next_dist.logits - T.logsumexp(next_dist.logits, dim=-1, keepdim=True)
+            q1_next = self.q1_target(next_states)
+            q2_next = self.q2_target(next_states)
+            # Soft Policy Evaluation 
+            min_q_next = T.min(q1_next, q2_next)
+            next_value = (next_probs * (min_q_next - self.alpha * next_log_probs)).sum(dim=-1, keepdim=True)
+            target = rewards + self.gamma * (1 - dones) * next_value
+
+        q1 = self.q1(states).gather(1, actions)
+        q2 = self.q2(states).gather(1, actions)
+
+        # Losses of both Q-functions
+        q1_loss = nn.MSELoss()(q1, target)
+        q2_loss = nn.MSELoss()(q2, target)
+
+        self.q1_opt.zero_grad()
+        q1_loss.backward()
+        self.q1_opt.step()
+        self.q2_opt.zero_grad()
+        q2_loss.backward()
+        self.q2_opt.step()
+
+        # Policy/Actor Objective
+        logits = self.policy(states)
+        dist = Categorical(logits=logits)
+        probs = dist.probs
+        log_probs = dist.logits - T.logsumexp(dist.logits, dim=-1, keepdim=True)
+        q1_policy = self.q1(states)
+        q2_policy = self.q2(states)
+        min_q_policy = T.min(q1_policy, q2_policy)
+        # Slightly different policy loss for discrete actions
+        policy_loss = (probs * (self.alpha * log_probs - min_q_policy)).sum(dim=-1).mean()
+
+        # Temperature to update Alpha
+        alpha_loss = -(self.log_alpha * (log_probs + self.target_entropy).detach()).mean()
+        self.alpha_opt.zero_grad()
+        alpha_loss.backward()
+        T.nn.utils.clip_grad_norm_([self.log_alpha], max_norm=1.0)  # Gradient clipping
+        self.alpha_opt.step()
+        self.alpha = self.log_alpha.exp().item()  
+
+        self.policy_opt.zero_grad()
+        policy_loss.backward()
+        T.nn.utils.clip_grad_norm_([self.log_alpha], max_norm=1.0)  # Gradient clipping
+        self.policy_opt.step()
+
+        # Target network update
+        for target_param, param in zip(self.q1_target.parameters(), self.q1.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+        for target_param, param in zip(self.q2_target.parameters(), self.q2.parameters()):
+            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+
+        return policy_loss.item()
+
+# Actor/Policy network 
+# Typical SAC Actor network is used to output a Gaussian distribution of a state
+# Here, we adapt it for discrete actions using a Categorical distribution, as the ATARI environment is discrete
+# The policy outputs logits for each discrete action.
+
+# From: https://ch.mathworks.com/help/reinforcement-learning/ug/soft-actor-critic-agents.html
+# The actor takes the current observation and generates a categorical distribution, in which each possible action is associated with a probability. 
+
+class CategoricalActor(nn.Module):
+    def __init__(self, obs_shape: tuple, action_dim: int, hidden: int):
+        super().__init__()
+        c, h, w = obs_shape
+        self.cnn = nn.Sequential(
+            nn.Conv2d(c, 16, kernel_size=8, stride=4),
+            nn.ReLU(),
+            nn.Conv2d(16, 32, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Flatten()
+        )
+        with T.no_grad():
+            cnn_output_dim = self.cnn(T.zeros(1, c, h, w)).shape[1]
+        self.fc = nn.Sequential(
+            nn.Linear(cnn_output_dim, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, action_dim)
+        )
+
+    def forward(self, state: T.Tensor):
+        if state.dim() == 3:
+            state = state.unsqueeze(0)
+        cnn_out = self.cnn(state)
+        logits = self.fc(cnn_out)
+        return logits
+
+# Q-network for discrete actions 
+class QNetwork(nn.Module):
+    def __init__(self, obs_shape: tuple, action_dim: int, hidden: int):
+        super().__init__()
+        c, h, w = obs_shape
+        self.cnn = nn.Sequential(
+            nn.Conv2d(c, 16, kernel_size=8, stride=4),
+            nn.ReLU(),
+            nn.Conv2d(16, 32, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Flatten()
+        )
+        with T.no_grad():
+            cnn_output_dim = self.cnn(T.zeros(1, c, h, w)).shape[1]
+        self.net = nn.Sequential(
+            nn.Linear(cnn_output_dim, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, action_dim)
+        )
+
+    def forward(self, state: T.Tensor):
+        if state.dim() == 3:
+            state = state.unsqueeze(0)
+        cnn_out = self.cnn(state)
+        return self.net(cnn_out)
+
+class Memory:
+    def __init__(self):
+        self.states, self.actions, self.rewards, self.dones, self.next_states = [], [], [], [], []
+    def store(self, s, a, r, d, ns):
+        self.states.append(np.asarray(s, dtype=np.float32))
+        self.actions.append(np.asarray(a, dtype=np.float32))
+        self.rewards.append(float(r))
+        self.dones.append(float(d))
+        self.next_states.append(np.asarray(ns, dtype=np.float32))
+    def clear(self):
+        self.states, self.actions, self.rewards, self.dones, self.next_states = [], [], [], [], []

Observation_norm_SAC/sac_model_cnn.py

@@ -0,0 +1,206 @@
+import ale_py
+import gymnasium as gym
+import sys
+import numpy as np
+from sac_helpers_cnn import *
+from gymnasium.spaces import Box
+import cv2
+import matplotlib.pyplot as plt
+
+def preprocess(obs):
+    obs = cv2.cvtColor(obs, cv2.COLOR_RGB2GRAY)
+    obs = cv2.resize(obs, (84, 84), interpolation=cv2.INTER_AREA)
+    obs = np.expand_dims(obs, axis=0) 
+    return obs.astype(np.float32) / 255.0
+"""
+def main() -> int:
+    episode = 0
+    total_return = 0
+    ep_return = 0
+    steps = 100
+    batches = 100
+    avg_returns = []
+    avg_losses = []
+
+    env = gym.make("ALE/Pacman-v5")
+    # Initialize CNN with a dummy observation (to get correct input shape)
+    obs, _ = env.reset()
+    dummy_obs_space = Box(low=0.0, high=1.0, shape=preprocess(obs).shape)
+
+    agent = Agent(
+        obs_space=dummy_obs_space,
+        action_space=env.action_space,
+        hidden=64,
+        gamma=0.99,
+        lr=3e-4,
+        alpha=0.2,
+        seed=70,
+        batch_size=32,
+        tau=0.005
+    )
+
+    try:
+        obs, info = env.reset(seed=42)
+        state = preprocess(obs)
+
+        for update in range(1, batches + 1):
+            batch_loss = []
+
+            for t in range(steps):
+                action = agent.choose_action(state)
+                next_obs, reward, terminated, truncated, info = env.step(action)
+                done = terminated or truncated
+                next_state = preprocess(next_obs)
+
+                agent.remember(state, action, reward, done, next_state)
+
+                ep_return += reward
+                state = next_state
+
+                if done:
+                    episode += 1
+                    total_return += ep_return
+                    print(f"Episode {episode} return: {ep_return:.2f}")
+                    ep_return = 0
+                    obs, info = env.reset()
+                    state = preprocess(obs)
+
+            loss = agent.vanilla_sac_update()
+            batch_loss.append(loss) 
+
+            avg_ret = (total_return / episode) if episode else 0
+            avg_returns.append(avg_ret)
+            avg_losses.append(np.mean(batch_loss)) 
+
+            print(f"Update {update}: episodes={episode}, avg_return={avg_ret:.2f}, avg_loss={np.mean(batch_loss):.4f}")
+            
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        return 1
+    finally:
+        avg = total_return / episode if episode else 0
+        print(f"\nEpisodes: {episode}, Avg return: {avg:.3f}")
+
+        # Plot learning curve
+        plt.figure(figsize=(10, 4))
+        plt.subplot(1, 2, 1)
+        plt.plot(avg_returns)
+        plt.xlabel("Update")
+        plt.ylabel("Average Return")
+        plt.title("SAC Learning Curve")
+        plt.grid()
+
+        plt.subplot(1, 2, 2)
+        plt.plot(avg_losses)
+        plt.xlabel("Update")
+        plt.ylabel("Average Loss")
+        plt.title("Average Loss Curve")
+        plt.grid()
+
+        plt.tight_layout()
+        plt.show()
+        env.close()
+
+    return 0
+"""
+
+def run_training(seed=42):
+    episode = 0
+    total_return = 0
+    ep_return = 0
+    steps = 100
+    batches = 100
+    avg_returns = []
+    avg_losses = []
+
+    env = gym.make("ALE/Pacman-v5")
+    obs, _ = env.reset()
+    dummy_obs_space = Box(low=0.0, high=1.0, shape=preprocess(obs).shape)
+
+    agent = Agent(
+        obs_space=dummy_obs_space,
+        action_space=env.action_space,
+        hidden=64,
+        gamma=0.99,
+        lr=3e-4,
+        alpha=0.2,
+        seed=seed,
+        batch_size=32,
+        tau=0.005
+    )
+
+    obs, info = env.reset(seed=seed)
+    state = preprocess(obs)
+
+    for update in range(1, batches + 1):
+        batch_loss = []
+        for t in range(steps):
+            action = agent.choose_action(state)
+            next_obs, reward, terminated, truncated, info = env.step(action)
+            done = terminated or truncated
+            next_state = preprocess(next_obs)
+
+            agent.remember(state, action, reward, done, next_state)
+
+            ep_return += reward
+            state = next_state
+
+            if done:
+                episode += 1
+                total_return += ep_return
+                ep_return = 0
+                obs, info = env.reset()
+                state = preprocess(obs)
+
+        #loss = agent.vanilla_sac_update()
+        loss = agent.update_reward_gradient_clipping()
+        batch_loss.append(loss)
+
+        avg_ret = (total_return / episode) if episode else 0
+        avg_returns.append(avg_ret)
+        avg_losses.append(np.mean(batch_loss))
+
+    env.close()
+    return avg_returns, avg_losses
+
+def main() -> int:
+    num_runs = 5
+    all_returns = []
+    all_losses = []
+
+    for run in range(num_runs):
+        print(f"Starting run {run+1}/{num_runs}")
+        avg_returns, avg_losses = run_training(seed=42 + run)
+        all_returns.append(avg_returns)
+        all_losses.append(avg_losses)
+
+    # Convert to numpy arrays for easy averaging
+    all_returns = np.array(all_returns)
+    all_losses = np.array(all_losses)
+
+    mean_returns = np.mean(all_returns, axis=0)
+    mean_losses = np.mean(all_losses, axis=0)
+
+    # Plot averaged learning curves
+    plt.figure(figsize=(10, 4))
+    plt.subplot(1, 2, 1)
+    plt.plot(mean_returns)
+    plt.xlabel("Update")
+    plt.ylabel("Average Return")
+    plt.title(f"SAC Learning Curve (avg over {num_runs} runs)")
+    plt.grid()
+
+    plt.subplot(1, 2, 2)
+    plt.plot(mean_losses)
+    plt.xlabel("Update")
+    plt.ylabel("Average Loss")
+    plt.title(f"Average Loss Curve (avg over {num_runs} runs)")
+    plt.grid()
+
+    plt.tight_layout()
+    plt.show()
+
+    return 0
+
+if __name__ == "__main__":
+    raise SystemExit(main())

SAC-2/sac-project/sac_helpers_cnn.py

@@ -11,6 +11,7 @@ class Agent:
             T.manual_seed(seed)
         
         # Use GPU if available
+
         self.device = T.device('cuda:0' if T.cuda.is_available() else 'cpu')
         self.action_dim = int(getattr(action_space, "n", action_space.n))  # Use .n for Discrete
         self.obs_shape = obs_space.shape