Soft Actor-Critic (SAC)算法

代码

1. import gym
2. import torch
3. import torch.nn as nn
4. import torch.optim as optim
5. import torch.nn.functional as F
6. import numpy as np
7. import pygame
9. # 定义 Actor 网络
10. class Actor(nn.Module):
11.     def __init__(self, state_dim, action_dim, max_action):
12.         super(Actor, self).__init__()
13.         self.fc1 = nn.Linear(state_dim, 256)
14.         self.fc2 = nn.Linear(256, 256)
15.         self.mu = nn.Linear(256, action_dim)
16.         self.log_std = nn.Linear(256, action_dim)
17.         self.max_action = max_action
19.     def forward(self, state):
20.         x = F.relu(self.fc1(state))
21.         x = F.relu(self.fc2(x))
22.         mu = self.mu(x)
23.         log_std = self.log_std(x)
24.         log_std = torch.clamp(log_std, -20, 2)
25.         std = torch.exp(log_std)
26.         return mu, std
28.     def sample(self, state):
29.         mu, std = self.forward(state)
30.         dist = torch.distributions.Normal(mu, std)
31.         action = dist.rsample()
32.         action = torch.tanh(action) * self.max_action
33.         log_prob = dist.log_prob(action).sum(axis=-1)
34.         log_prob -= (2 * (np.log(2) - action - F.softplus(-2 * action))).sum(axis=-1)
35.         return action, log_prob
37. # 定义 Critic 网络
38. class Critic(nn.Module):
39.     def __init__(self, state_dim, action_dim):
40.         super(Critic, self).__init__()
41.         self.fc1 = nn.Linear(state_dim + action_dim, 256)
42.         self.fc2 = nn.Linear(256, 256)
43.         self.fc3 = nn.Linear(256, 1)
45.     def forward(self, state, action):
46.         x = torch.cat([state, action], 1)
47.         x = F.relu(self.fc1(x))
48.         x = F.relu(self.fc2(x))
49.         x = self.fc3(x)
50.         return x
52. # SAC 算法
53. class SAC:
54.     def __init__(self, state_dim, action_dim, max_action):
55.         self.actor = Actor(state_dim, action_dim, max_action)
56.         self.critic1 = Critic(state_dim, action_dim)
57.         self.critic2 = Critic(state_dim, action_dim)
58.         self.target_critic1 = Critic(state_dim, action_dim)
59.         self.target_critic2 = Critic(state_dim, action_dim)
61.         self.target_critic1.load_state_dict(self.critic1.state_dict())
62.         self.target_critic2.load_state_dict(self.critic2.state_dict())
64.         self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=3e-4)
65.         self.critic1_optimizer = optim.Adam(self.critic1.parameters(), lr=3e-4)
66.         self.critic2_optimizer = optim.Adam(self.critic2.parameters(), lr=3e-4)
68.         self.log_alpha = torch.tensor(np.log(0.1), requires_grad=True)
69.         self.alpha_optimizer = optim.Adam([self.log_alpha], lr=3e-4)
71.         self.gamma = 0.99
72.         self.tau = 0.005
74.     def select_action(self, state):
75.         state = torch.FloatTensor(state.reshape(1, -1))  # 确保 state 是 (1, state_dim) 的形状
76.         action, _ = self.actor.sample(state)
77.         return action.cpu().data.numpy().flatten()
79.     def update(self, replay_buffer, batch_size=256):
80.         state, action, next_state, reward, done = replay_buffer.sample(batch_size)
82.         state = torch.FloatTensor(state)
83.         action = torch.FloatTensor(action)
84.         next_state = torch.FloatTensor(next_state)
85.         reward = torch.FloatTensor(reward).unsqueeze(1)
86.         done = torch.FloatTensor(done).unsqueeze(1)
88.         with torch.no_grad():
89.             next_action, next_log_prob = self.actor.sample(next_state)
90.             target_q1 = self.target_critic1(next_state, next_action)
91.             target_q2 = self.target_critic2(next_state, next_action)
92.             target_q = torch.min(target_q1, target_q2) - self.log_alpha.exp() * next_log_prob
93.             target_q = reward + (1 - done) * self.gamma * target_q
95.         current_q1 = self.critic1(state, action)
96.         current_q2 = self.critic2(state, action)
97.         critic1_loss = F.mse_loss(current_q1, target_q)
98.         critic2_loss = F.mse_loss(current_q2, target_q)
100.         self.critic1_optimizer.zero_grad()
101.         critic1_loss.backward()
102.         self.critic1_optimizer.step()
104.         self.critic2_optimizer.zero_grad()
105.         critic2_loss.backward()
106.         self.critic2_optimizer.step()
108.         action_new, log_prob = self.actor.sample(state)
109.         q1_new = self.critic1(state, action_new)
110.         q2_new = self.critic2(state, action_new)
111.         q_new = torch.min(q1_new, q2_new)
112.         actor_loss = (self.log_alpha.exp() * log_prob - q_new).mean()
114.         self.actor_optimizer.zero_grad()
115.         actor_loss.backward()
116.         self.actor_optimizer.step()
118.         alpha_loss = -(self.log_alpha * (log_prob + 1).detach()).mean()
119.         self.alpha_optimizer.zero_grad()
120.         alpha_loss.backward()
121.         self.alpha_optimizer.step()
123.         for param, target_param in zip(self.critic1.parameters(), self.target_critic1.parameters()):
124.             target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
126.         for param, target_param in zip(self.critic2.parameters(), self.target_critic2.parameters()):
127.             target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
129. # 简单的 Replay Buffer
130. class ReplayBuffer:
131.     def __init__(self, max_size=1e6):
132.         self.buffer = []
133.         self.max_size = int(max_size)  # 将 max_size 转换为整数
134.         self.ptr = 0
136.     def add(self, state, action, next_state, reward, done):
137.         if len(self.buffer) < self.max_size:
138.             self.buffer.append(None)
139.         self.buffer[self.ptr] = (state, action, next_state, reward, done)
140.         self.ptr = (self.ptr + 1) % self.max_size
142.     def sample(self, batch_size):
143.         indices = np.random.randint(0, len(self.buffer), batch_size)
144.         states, actions, next_states, rewards, dones = [], [], [], [], []
145.         for idx in indices:
146.             state, action, next_state, reward, done = self.buffer[idx]
147.             states.append(state)
148.             actions.append(action)
149.             next_states.append(next_state)
150.             rewards.append(reward)
151.             dones.append(done)
152.         return np.array(states), np.array(actions), np.array(next_states), np.array(rewards), np.array(dones)
154. # 训练 SAC 算法
155. env = gym.make('Pendulum-v1')
156. state_dim = env.observation_space.shape[0]
157. action_dim = env.action_space.shape[0]
158. max_action = float(env.action_space.high[0])
160. sac = SAC(state_dim, action_dim, max_action)
161. replay_buffer = ReplayBuffer()
163. max_episodes = 1000
164. batch_size = 256
166. for episode in range(max_episodes):
167.     state = env.reset()
168.     if isinstance(state, tuple):  # 如果返回的是元组，提取状态
169.         state = state[0]
170.     episode_reward = 0
171.     done = False
172.     while not done:
173.         env.render()
174.         action = sac.select_action(state)
175.         next_state, reward, done, info = env.step(action)
176.         replay_buffer.add(state, action, next_state, reward, done)
177.         state = next_state
178.         episode_reward += reward
180.         if len(replay_buffer.buffer) > batch_size:
181.             sac.update(replay_buffer, batch_size)
183.     print(f"Episode {episode + 1}, Reward: {episode_reward}")
185. env.close()

复制代码

简介

Soft Actor-Critic (SAC) 是一种基于最大熵（Maximum Entropy）的深度强化学习算法，专为一连动作空间设计。它联合了 Actor-Critic 框架和熵正则化（Entropy Regularization），在探索与利用之间取得了良好的平衡。
图片

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。