60 lines
2.5 KiB
Python
60 lines
2.5 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
class ReplayBuffer(object):
|
|
def __init__(self, state_dim, action_dim, max_size=int(1e6)):
|
|
"""
|
|
初始化经验回放池
|
|
使用预分配的 Numpy 数组来提升存储和采样效率
|
|
"""
|
|
self.max_size = max_size
|
|
self.ptr = 0 # 当前写入的指针位置
|
|
self.size = 0 # 当前池子里的有效数据量
|
|
|
|
# 预先分配内存，避免动态扩张带来性能开销
|
|
self.state = np.zeros((max_size, state_dim), dtype=np.float32)
|
|
self.action = np.zeros((max_size, action_dim), dtype=np.float32)
|
|
self.reward = np.zeros((max_size, 1), dtype=np.float32)
|
|
self.next_state = np.zeros((max_size, state_dim), dtype=np.float32)
|
|
|
|
# 记录该状态是否是回合的结束 (1.0 表示结束，0.0 表示未结束)
|
|
self.not_done = np.zeros((max_size, 1), dtype=np.float32)
|
|
|
|
# 自动检测 GPU
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
def add(self, state, action, reward, next_state, done):
|
|
"""
|
|
向回放池中添加一条新的转移数据 (Transition)
|
|
"""
|
|
# 将数据写入指针当前所在的位置
|
|
self.state[self.ptr] = state
|
|
self.action[self.ptr] = action
|
|
self.reward[self.ptr] = reward
|
|
self.next_state[self.ptr] = next_state
|
|
|
|
# 我们存储 1 - done，这样在贝尔曼方程更新时直接相乘即可：
|
|
# Q_target = r + gamma * V * not_done
|
|
self.not_done[self.ptr] = 1. - done
|
|
|
|
# 移动指针，如果达到了最大容量，就回到开头覆盖最老的数据（环形结构）
|
|
self.ptr = (self.ptr + 1) % self.max_size
|
|
# 更新当前有效数据量
|
|
self.size = min(self.size + 1, self.max_size)
|
|
|
|
def sample(self, batch_size):
|
|
"""
|
|
随机采样一个 batch 的数据，并直接转换为 PyTorch Tensor 放到 GPU/CPU 上
|
|
"""
|
|
# 在 0 到当前有效数据量之间，随机生成 batch_size 个索引
|
|
ind = np.random.randint(0, self.size, size=batch_size)
|
|
|
|
# 提取数据并转为 Tensor
|
|
return (
|
|
torch.FloatTensor(self.state[ind]).to(self.device),
|
|
torch.FloatTensor(self.action[ind]).to(self.device),
|
|
torch.FloatTensor(self.reward[ind]).to(self.device),
|
|
torch.FloatTensor(self.next_state[ind]).to(self.device),
|
|
torch.FloatTensor(self.not_done[ind]).to(self.device)
|
|
)
|