A Comprehensive Overview of Q-Learning and Actor-Critic Methods

1
import numpy as np
2

3
# Suppose Q is a table (2D array: Q[state, action])
4
def q_learning_update(Q, s, a, r, s_next, alpha, gamma):
5
    # Q[s, a] = Q[s, a] + alpha * [r + gamma * max(Q[s_next, :]) - Q[s, a]]
6
    Q[s, a] += alpha * (r + gamma * np.max(Q[s_next]) - Q[s, a])

1
import torch
2
import torch.nn as nn
3
import torch.optim as optim
4

5
class QNetwork(nn.Module):
6
    def __init__(self, state_dim, action_dim, hidden_dim=64):
7
        super(QNetwork, self).__init__()
8
        self.fc1 = nn.Linear(state_dim, hidden_dim)
9
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
10
        self.fc3 = nn.Linear(hidden_dim, action_dim)  # Q(s) -> R^action_dim
11

12
    def forward(self, s):
13
        x = torch.relu(self.fc1(s))
14
        x = torch.relu(self.fc2(x))
15
        return self.fc3(x)
16

17
def dqn_loss(q_network, target_network, batch, gamma):
18
    # batch includes (s, a, r, s_next, done)
19
    s, a, r, s_next, done = batch
20

21
    # Q(s, a)
22
    q_values = q_network(s).gather(1, a.unsqueeze(1)).squeeze(1)
23

24
    # target = r + gamma * max_a' Q^- (s_next, a') if not done
25
    with torch.no_grad():
26
        q_next = target_network(s_next).max(1)[0]
27
        q_target = r + gamma * q_next * (1 - done.float())
28

29
    loss = nn.MSELoss()(q_values, q_target)
30
    return loss

1
import torch
2

3
def policy_gradient_loss(log_probs, returns, baselines):
4
    # log_probs: tensor of log π(a|s) for each step
5
    # returns: tensor of G_t
6
    # baselines: tensor of b(s), often V(s)
7

8
    advantage = returns - baselines
9
    # actor loss = - E[ log pi(a|s) * advantage ]
10
    loss = -(log_probs * advantage).mean()
11
    return loss

1
import torch.nn as nn
2

3
def value_loss(value_net, states, returns):
4
    # value_net(s) -> scalar V_phi(s)
5
    v = value_net(states)
6
    loss = nn.MSELoss()(v.squeeze(), returns)
7
    return loss

1
class MLPActor(nn.Module):
2
    def __init__(self, state_dim, action_dim, hidden_dim=64, discrete=False):
3
        super(MLPActor, self).__init__()
4
        self.fc1 = nn.Linear(state_dim, hidden_dim)
5
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
6
        # Discrete: outputs action_dim logits
7
        # Continuous: outputs mean, possibly log std
8
        self.discrete = discrete
9
        if discrete:
10
            self.fc_out = nn.Linear(hidden_dim, action_dim)
11
        else:
12
            # For continuous, let's say we output mean only
13
            self.fc_mean = nn.Linear(hidden_dim, action_dim)
14
            # could also have self.log_std = nn.Parameter(...)
15

16
    def forward(self, s):
17
        x = torch.relu(self.fc1(s))
18
        x = torch.relu(self.fc2(x))
19
        if self.discrete:
20
            logits = self.fc_out(x)
21
            return logits  # use softmax outside
22
        else:
23
            mean = self.fc_mean(x)
24
            return mean

1
def ddpg_critic_loss(q_net, target_q_net, batch, gamma, actor, target_actor):
2
    s, a, r, s_next, done = batch
3
    q_vals = q_net(s, a)
4
    with torch.no_grad():
5
        a_next = target_actor(s_next)
6
        q_next = target_q_net(s_next, a_next)
7
        q_target = r + gamma * q_next * (1 - done)
8
    loss = nn.MSELoss()(q_vals, q_target)
9
    return loss

1
def ddpg_actor_loss(q_net, states, actor):
2
    # actor(s) -> a
3
    a = actor(states)
4
    # compute d/d(a) of Q(s,a), then chain rule
5
    q_vals = q_net(states, a)
6
    # we want to maximize Q, so minimize negative
7
    loss = -q_vals.mean()
8
    return loss

1
def td3_critic_loss(q1, q2, q1_target, q2_target, batch, gamma, actor_target):
2
    s, a, r, s_next, done = batch
3

4
    with torch.no_grad():
5
        a_next = actor_target(s_next)
6
        q1_next = q1_target(s_next, a_next)
7
        q2_next = q2_target(s_next, a_next)
8
        q_next_min = torch.min(q1_next, q2_next)
9
        q_target = r + gamma * q_next_min * (1 - done)
10

11
    loss1 = nn.MSELoss()(q1(s, a), q_target)
12
    loss2 = nn.MSELoss()(q2(s, a), q_target)
13
    return loss1 + loss2

1
def ppo_clip_loss(pi_new, pi_old, actions, advantages, epsilon=0.2):
2
    # pi_new, pi_old: probability of actions under new/old policy
3
    ratio = pi_new / (pi_old + 1e-8)
4
    unclipped = ratio * advantages
5
    clipped = torch.clamp(ratio, 1-epsilon, 1+epsilon) * advantages
6
    loss = -torch.min(unclipped, clipped).mean()
7
    return loss

1
def sac_critic_loss(q1, q2, q1_target, q2_target, batch, alpha, gamma, actor_target):
2
    s, a, r, s_next, done = batch
3
    with torch.no_grad():
4
        # sample new action from actor_target
5
        a_next = actor_target(s_next)
6
        # log pi(a_next|s_next)
7
        logp_a_next = actor_target.log_prob(s_next, a_next)
8
        q1_next = q1_target(s_next, a_next)
9
        q2_next = q2_target(s_next, a_next)
10
        q_next_min = torch.min(q1_next, q2_next)
11
        q_target = r + gamma * (q_next_min - alpha * logp_a_next) * (1 - done)
12

13
    loss1 = nn.MSELoss()(q1(s, a), q_target)
14
    loss2 = nn.MSELoss()(q2(s, a), q_target)
15
    return loss1 + loss2

1
def a3c_loss(policy_net, value_net, states, actions, returns):
2
    # policy_net -> log pi(a|s), value_net -> V(s)
3
    log_probs = policy_net.log_prob(states, actions)
4
    values = value_net(states).squeeze()
5
    advantage = returns - values
6

7
    actor_loss = - (log_probs * advantage).mean()
8
    critic_loss = advantage.pow(2).mean()
9
    return actor_loss + critic_loss

1
# -*- coding: utf-8 -*-
2
"""
3
Refactored A3C Example:
4
  - Renamed classes, functions, and variables for clarity.
5
  - Maintains the original multi-process A3C logic with shared parameters.
6

7
Author: <Your Name>
8
"""
9

10
import gymnasium as gym
11
import torch
12
import torch.nn as nn
13
import torch.nn.functional as F
14
import torch.multiprocessing as mp
15
from torch.distributions import Categorical
16

17

18
##############################################################################
19
# 1) UTILITY: A safer environment "step" handling older/newer Gym returns.
20
##############################################################################
21
def safe_step_env(env, action):
22
    """
23
    Step the environment while handling Gym/Gymnasium return formats:
24
      - Some return (obs, reward, done, info)
25
      - Gymnasium can return (obs, reward, done, truncated, info)
26
    """
27
    results = env.step(action)
28
    if len(results) == 5:
29
        next_obs, reward, done, truncated, info = results
30
        done = done or truncated
31
    else:
32
        next_obs, reward, done, info = results
33
    return next_obs, reward, done, info
34

35

36
##############################################################################
37
# 2) SHARED OPTIMIZER
38
##############################################################################
39
class SharedAdam(torch.optim.Adam):
40
    """
41
    A custom Adam optimizer that uses shared memory for the moving averages,
42
    allowing multiple processes to update a shared global set of parameters.
43
    """
44

45
    def __init__(self, parameters, lr=1e-3, betas=(0.9, 0.99), eps=1e-8,
46
                 weight_decay=0):
47
        super(SharedAdam, self).__init__(parameters, lr=lr, betas=betas,
48
                                         eps=eps, weight_decay=weight_decay)
49
        for group in self.param_groups:
50
            for param in group['params']:
51
                state = self.state[param]
52
                state['step'] = 0
53
                state['exp_avg'] = torch.zeros_like(param.data)
54
                state['exp_avg_sq'] = torch.zeros_like(param.data)
55
                # Share memory so that parallel processes can update
56
                # the optimizer states
57
                state['exp_avg'].share_memory_()
58
                state['exp_avg_sq'].share_memory_()
59

60

61
##############################################################################
62
# 3) A3C NETWORK (Actor + Critic)
63
##############################################################################
64
class A3CNetwork(nn.Module):
65
    """
66
    An Actor-Critic network that outputs:
67
      - A policy distribution (pi) for selecting actions.
68
      - A value function (v) for estimating state value.
69
    """
70

71
    def __init__(self, state_dim, num_actions, gamma=0.99):
72
        super(A3CNetwork, self).__init__()
73
        self.gamma = gamma
74

75
        self.actor_hidden = nn.Linear(*state_dim, 128)
76
        self.critic_hidden = nn.Linear(*state_dim, 128)
77

78
        self.actor_out = nn.Linear(128, num_actions)
79
        self.critic_out = nn.Linear(128, 1)
80

81
        # We'll store trajectories in memory buffers between updates
82
        self.memory_states = []
83
        self.memory_actions = []
84
        self.memory_rewards = []
85

86
    def store_experience(self, state, action, reward):
87
        """Cache a transition (state, action, reward)."""
88
        self.memory_states.append(state)
89
        self.memory_actions.append(action)
90
        self.memory_rewards.append(reward)
91

92
    def reset_memory(self):
93
        """Clear trajectory buffers."""
94
        self.memory_states = []
95
        self.memory_actions = []
96
        self.memory_rewards = []
97

98
    def forward(self, state):
99
        """
100
        Forward pass:
101
          state: a torch.Tensor of shape [batch_size, *state_dim].
102
        Returns:
103
          - policy_logits (for the actor)
104
          - state_value (scalar per batch item, for the critic)
105
        """
106
        actor_hidden_out = F.relu(self.actor_hidden(state))
107
        critic_hidden_out = F.relu(self.critic_hidden(state))
108

109
        policy_logits = self.actor_out(actor_hidden_out)
110
        state_value = self.critic_out(critic_hidden_out)
111
        return policy_logits, state_value
112

113
    def compute_returns(self, done):
114
        """
115
        Compute discounted returns for each step in the trajectory.
116
        If the episode ended (done=True), final state's value is 0.
117
        Otherwise, we bootstrap from the last state's value.
118
        """
119
        states_t = torch.tensor(self.memory_states, dtype=torch.float)
120
        _, values_t = self.forward(states_t)
121
        values_t = values_t.squeeze(-1)  # shape [trajectory_len]
122

123
        # If done, final state's value is 0; else we take the last state's value
124
        final_value = 0.0 if done else values_t[-1].item()
125

126
        returns = []
127
        discounted_return = final_value
128
        for reward in reversed(self.memory_rewards):
129
            discounted_return = reward + self.gamma * discounted_return
130
            returns.append(discounted_return)
131
        returns.reverse()
132
        return torch.tensor(returns, dtype=torch.float)
133

134
    def compute_loss(self, done):
135
        """
136
        Compute combined Actor + Critic loss over the stored trajectory.
137
        """
138
        states_t = torch.tensor(self.memory_states, dtype=torch.float)
139
        actions_t = torch.tensor(self.memory_actions, dtype=torch.long)
140
        returns_t = self.compute_returns(done)
141

142
        policy_logits, values_t = self.forward(states_t)
143
        values_t = values_t.squeeze(-1)  # shape [trajectory_len]
144

145
        # Critic loss (Mean Squared Error)
146
        critic_loss = (returns_t - values_t) ** 2
147

148
        # Actor loss (REINFORCE with advantage)
149
        probabilities = F.softmax(policy_logits, dim=1)
150
        dist = Categorical(probabilities)
151
        log_probs = dist.log_prob(actions_t)
152

153
        advantages = returns_t - values_t.detach()  # no grad wrt. values for advantage
154
        actor_loss = -log_probs * advantages
155

156
        total_loss = (critic_loss + actor_loss).mean()
157
        return total_loss
158

159
    def select_action(self, observation):
160
        """
161
        Sample an action according to the policy's distribution.
162
        observation: raw state from the environment (np array or float list).
163
        """
164
        observation_t = torch.tensor([observation], dtype=torch.float)
165
        policy_logits, _ = self.forward(observation_t)
166
        probabilities = F.softmax(policy_logits, dim=1)
167
        dist = Categorical(probabilities)
168
        action = dist.sample().item()
169
        return action
170

171

172
##############################################################################
173
# 4) WORKER AGENT (One process per agent)
174
##############################################################################
175
class A3CWorker(mp.Process):
176
    """
177
    Each worker interacts with an environment instance, accumulates experience,
178
    and updates the global A3C network's parameters.
179
    """
180

181
    def __init__(self,
182
                 global_network,
183
                 global_optimizer,
184
                 state_dim,
185
                 num_actions,
186
                 gamma,
187
                 lr,
188
                 worker_id,
189
                 global_episode_counter,
190
                 env_id,
191
                 max_episodes,
192
                 update_interval):
193
        super(A3CWorker, self).__init__()
194
        self.local_network = A3CNetwork(state_dim, num_actions, gamma)
195
        self.global_network = global_network
196
        self.global_optimizer = global_optimizer
197

198
        self.worker_name = f"worker_{worker_id:02d}"
199
        self.episode_counter = global_episode_counter
200

201
        self.env = gym.make(env_id)
202
        self.max_episodes = max_episodes
203
        self.update_interval = update_interval
204

205
    def run(self):
206
        step_count = 1
207
        while self.episode_counter.value < self.max_episodes:
208
            state, _info = self._reset_env()  # handle Gymnasium reset
209
            done = False
210
            episode_return = 0.0
211

212
            self.local_network.reset_memory()
213

214
            while not done:
215
                action = self.local_network.select_action(state)
216
                next_state, reward, done, info = safe_step_env(self.env, action)
217

218
                episode_return += reward
219
                self.local_network.store_experience(state, action, reward)
220

221
                # Update global network after 'update_interval' steps or on episode done
222
                if step_count % self.update_interval == 0 or done:
223
                    loss = self.local_network.compute_loss(done)
224

225
                    self.global_optimizer.zero_grad()
226
                    loss.backward()
227

228
                    # Copy grads from local to global
229
                    for local_param, global_param in zip(
230
                            self.local_network.parameters(),
231
                            self.global_network.parameters()):
232
                        global_param._grad = local_param.grad
233

234
                    self.global_optimizer.step()
235

236
                    # Sync local network with updated global parameters
237
                    self.local_network.load_state_dict(self.global_network.state_dict())
238
                    self.local_network.reset_memory()
239

240
                step_count += 1
241
                state = next_state
242

243
            # Increment global episode counter
244
            with self.episode_counter.get_lock():
245
                self.episode_counter.value += 1
246

247
            print(f"{self.worker_name} | Episode: {self.episode_counter.value} "
248
                  f"| Return: {episode_return:.1f}")
249

250
    def _reset_env(self):
251
        """
252
        Reset environment (Gym/Gymnasium). Handles new reset() returning (obs, info).
253
        """
254
        initial_obs = self.env.reset()
255
        if isinstance(initial_obs, tuple) and len(initial_obs) == 2:
256
            obs, info = initial_obs
257
        else:
258
            obs = initial_obs
259
            info = {}
260
        return obs, info
261

262

263
##############################################################################
264
# 5) WATCH A TRAINED AGENT
265
##############################################################################
266
def watch_agent(global_network, env_id="CartPole-v1", episodes_to_watch=5):
267
    """
268
    Renders a few episodes using the global A3C network's parameters.
269
    """
270
    env = gym.make(env_id, render_mode="human")
271

272
    # Local copy for inference
273
    local_network = A3CNetwork([4], 2)  # For CartPole: state_dim=[4], num_actions=2
274
    local_network.load_state_dict(global_network.state_dict())
275

276
    for ep in range(episodes_to_watch):
277
        state, _info = env.reset()
278
        done = False
279
        episode_return = 0.0
280
        while not done:
281
            # For older Gym versions, you might do env.render() here
282
            action = local_network.select_action(state)
283
            state, reward, done, info = safe_step_env(env, action)
284
            episode_return += reward
285

286
        print(f"Watch Episode {ep + 1}, Return: {episode_return:.1f}")
287

288
    env.close()
289

290

291
##############################################################################
292
# 6) MAIN: TRAINING LOGIC
293
##############################################################################
294
if __name__ == '__main__':
295

296
    LEARNING_RATE = 1e-4
297
    ENV_ID = "CartPole-v1"
298
    env = gym.make(ENV_ID)
299

300
    # Observation space shape => typically something like (4,)
301
    state_dim = env.observation_space.shape
302

303
    # Action space => for Discrete(n), .n is the number of possible actions
304
    num_actions = env.action_space.n
305
    MAX_EPISODES = 3000
306
    UPDATE_INTERVAL = 500
307

308
    # Create global (shared) A3C network
309
    shared_network = A3CNetwork(state_dim, num_actions)
310
    shared_network.share_memory()
311

312
    # Create shared optimizer
313
    shared_optimizer = SharedAdam(shared_network.parameters(),
314
                                  lr=LEARNING_RATE,
315
                                  betas=(0.92, 0.999))
316

317
    global_episode_counter = mp.Value('i', 0)
318

319
    # Spawn worker processes
320
    num_cpus = mp.cpu_count()
321
    workers = []
322
    for cpu_id in range(num_cpus):
323
        worker = A3CWorker(
324
            global_network=shared_network,
325
            global_optimizer=shared_optimizer,
326
            state_dim=state_dim,
327
            num_actions=num_actions,
328
            gamma=0.99,
329
            lr=LEARNING_RATE,
330
            worker_id=cpu_id,
331
            global_episode_counter=global_episode_counter,
332
            env_id=ENV_ID,
333
            max_episodes=MAX_EPISODES,
334
            update_interval=UPDATE_INTERVAL
335
        )
336
        workers.append(worker)
337

338
    # Start and join each worker
339
    for w in workers:
340
        w.start()
341
    for w in workers:
342
        w.join()
343

344
    print("Training complete. Now let's watch the agent in action!")
345
    watch_agent(shared_network, env_id=ENV_ID, episodes_to_watch=5000)