添加 PPO 算法实现及相关配置,更新训练入口以支持 SAC 和 PPO 模式
This commit is contained in:
@@ -0,0 +1,104 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
class RolloutBuffer:
|
||||
"""
|
||||
On-policy 滚动缓冲区,用于 PPO 算法。
|
||||
|
||||
每次 collect() 收集 T 步数据后,调用 compute_returns_and_advantages()
|
||||
计算 GAE 优势估计,然后通过 get_batches() 将数据切分为 mini-batch
|
||||
供 K 轮 epoch 使用,最后 clear() 清空等待下一轮收集。
|
||||
"""
|
||||
|
||||
def __init__(self, state_dim, action_dim, steps_per_update, gamma, gae_lambda, device):
|
||||
self.steps = steps_per_update
|
||||
self.gamma = gamma
|
||||
self.gae_lambda = gae_lambda
|
||||
self.device = device
|
||||
|
||||
# 预分配存储空间
|
||||
self.states = np.zeros((steps_per_update, state_dim), dtype=np.float32)
|
||||
self.actions_unbounded = np.zeros((steps_per_update, action_dim), dtype=np.float32) # clamp 之前的高斯采样值
|
||||
self.log_probs = np.zeros((steps_per_update, 1), dtype=np.float32)
|
||||
self.rewards = np.zeros((steps_per_update, 1), dtype=np.float32)
|
||||
self.dones = np.zeros((steps_per_update, 1), dtype=np.float32)
|
||||
self.values = np.zeros((steps_per_update, 1), dtype=np.float32)
|
||||
|
||||
# 计算后填充
|
||||
self.returns = np.zeros((steps_per_update, 1), dtype=np.float32)
|
||||
self.advantages = np.zeros((steps_per_update, 1), dtype=np.float32)
|
||||
|
||||
self.ptr = 0
|
||||
self.full = False
|
||||
|
||||
def add(self, state, action_unbounded, log_prob, reward, done, value):
|
||||
"""
|
||||
向缓冲区写入一步数据。
|
||||
action_unbounded: 未裁剪的高斯采样值(形状 [action_dim])
|
||||
log_prob: 该步的 log π(a|s)(标量)
|
||||
value: V(s) 的估计值(标量)
|
||||
"""
|
||||
idx = self.ptr
|
||||
self.states[idx] = state
|
||||
self.actions_unbounded[idx] = action_unbounded
|
||||
self.log_probs[idx] = log_prob
|
||||
self.rewards[idx] = reward
|
||||
self.dones[idx] = done
|
||||
self.values[idx] = value
|
||||
self.ptr += 1
|
||||
if self.ptr >= self.steps:
|
||||
self.full = True
|
||||
|
||||
def compute_returns_and_advantages(self, last_value):
|
||||
"""
|
||||
反向遍历轨迹,计算 GAE 优势估计(论文公式11/12)。
|
||||
只对实际填充的 self.ptr 步数据计算,避免无效数据参与。
|
||||
|
||||
Args:
|
||||
last_value: V(s_{T+1}),下一个状态的价值估计(若 episode 结束则为 0)
|
||||
"""
|
||||
n = self.ptr # 实际有效数据条数
|
||||
last_gae = 0.0
|
||||
for t in reversed(range(n)):
|
||||
if t == n - 1:
|
||||
next_non_terminal = 1.0 - self.dones[t]
|
||||
next_value = last_value
|
||||
else:
|
||||
next_non_terminal = 1.0 - self.dones[t]
|
||||
next_value = self.values[t + 1]
|
||||
|
||||
delta = self.rewards[t] + self.gamma * next_value * next_non_terminal - self.values[t]
|
||||
last_gae = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae
|
||||
self.advantages[t] = last_gae
|
||||
|
||||
# 回报 G_t = Â_t + V(s_t)
|
||||
self.returns[:n] = self.advantages[:n] + self.values[:n]
|
||||
|
||||
# 优势归一化
|
||||
adv = self.advantages[:n]
|
||||
self.advantages[:n] = (adv - adv.mean()) / (adv.std() + 1e-8)
|
||||
|
||||
def get_batches(self, batch_size):
|
||||
"""
|
||||
将实际填充的数据随机打乱后按 batch_size 切片,生成 mini-batch。
|
||||
|
||||
Yields:
|
||||
(states, actions_unbounded, log_probs, returns, advantages) — 均为 Tensor
|
||||
"""
|
||||
n = self.ptr
|
||||
indices = np.random.permutation(n)
|
||||
for start in range(0, n, batch_size):
|
||||
idx = indices[start: start + batch_size]
|
||||
yield (
|
||||
torch.FloatTensor(self.states[idx]).to(self.device),
|
||||
torch.FloatTensor(self.actions_unbounded[idx]).to(self.device),
|
||||
torch.FloatTensor(self.log_probs[idx]).to(self.device),
|
||||
torch.FloatTensor(self.returns[idx]).to(self.device),
|
||||
torch.FloatTensor(self.advantages[idx]).to(self.device),
|
||||
)
|
||||
|
||||
def clear(self):
|
||||
"""清空缓冲区,为下一轮收集做准备。"""
|
||||
self.ptr = 0
|
||||
self.full = False
|
||||
Reference in New Issue
Block a user