AWR强化学习从入门到实战：基于PyTorch的智能体训练指南

1次阅读

共计 4284 个字符，预计需要花费 11 分钟才能阅读完成。

强化学习（Reinforcement Learning）近年来在多个领域取得了显著成果，而 Advantage-Weighted Regression (AWR) 作为一种 off-policy 算法，因其稳定性和高效性受到广泛关注。相比于 PPO（Proximal Policy Optimization）和 SAC（Soft Actor-Critic），AWR 在策略优化上更加简洁，无需复杂的价值函数近似或熵正则化，特别适合初学者理解和实现。

AWR 的核心思想是通过加权回归更新策略，权重由优势函数（advantage function）决定。这种方法避免了传统策略梯度算法的高方差问题，同时保持了 off-policy 学习的灵活性。接下来，我们将从工程实现的角度，逐步解析 AWR 算法的关键模块。

经验回放（Experience Replay）是 off-policy 学习的核心组件，AWR 通过优先级采样（Priority Sampling）进一步提升数据利用率。以下是基于 PyTorch 的实现代码：

import torch
from collections import deque
import numpy as np

class PrioritizedReplayBuffer:
    def __init__(self, capacity, alpha=0.6):
        self.capacity = capacity
        self.alpha = alpha  # 优先级指数
        self.buffer = deque(maxlen=capacity)
        self.priorities = deque(maxlen=capacity)

    def add(self, transition):
        """
        添加经验到缓冲区，初始优先级设为最大值
        transition: (state, action, reward, next_state, done)
        """
        max_priority = max(self.priorities) if self.priorities else 1.0
        self.buffer.append(transition)
        self.priorities.append(max_priority)

    def sample(self, batch_size, beta=0.4):
        """
        基于优先级采样 batch，beta 用于调节重要性采样权重
        返回: (states, actions, rewards, next_states, dones), indices, weights
        """
        priorities = np.array(self.priorities)
        probs = priorities ** self.alpha
        probs /= probs.sum()

        indices = np.random.choice(len(self.buffer), batch_size, p=probs)
        samples = [self.buffer[idx] for idx in indices]

        # 重要性采样权重
        weights = (len(self.buffer) * probs[indices]) ** (-beta)
        weights /= weights.max()

        states = torch.FloatTensor(np.vstack([x[0] for x in samples]))
        actions = torch.FloatTensor(np.vstack([x[1] for x in samples]))
        rewards = torch.FloatTensor(np.vstack([x[2] for x in samples]))
        next_states = torch.FloatTensor(np.vstack([x[3] for x in samples]))
        dones = torch.FloatTensor(np.vstack([x[4] for x in samples]))

        return (states, actions, rewards, next_states, dones), indices, torch.FloatTensor(weights)

    def update_priorities(self, indices, new_priorities):
        """更新采样经验的优先级"""
        for idx, priority in zip(indices, new_priorities):
            self.priorities[idx] = priority

优势函数 $A(s,a) = Q(s,a) – V(s)$ 衡量动作的相对好坏。实践中，我们常用 GAE（Generalized Advantage Estimation）计算优势值：

def compute_advantages(rewards, values, dones, gamma=0.99, lam=0.95):
    """
    使用 GAE 计算优势值
    rewards: 轨迹的奖励序列 [T]
    values: 状态价值估计 [T]
    dones: 终止标志 [T]
    gamma: 折扣因子
    lam: GAE 参数
    """
    advantages = torch.zeros_like(rewards)
    last_advantage = 0

    for t in reversed(range(len(rewards))):
        if dones[t]:
            delta = rewards[t] - values[t]
            last_advantage = 0
        else:
            delta = rewards[t] + gamma * values[t+1] - values[t]
        advantages[t] = delta + gamma * lam * last_advantage
        last_advantage = advantages[t]

    # 标准化优势值
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
    return advantages

AWR 通过 KL 散度约束策略更新幅度，保证训练稳定性：

class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_size=256):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.mean = nn.Linear(hidden_size, action_dim)
        self.log_std = nn.Parameter(torch.zeros(action_dim))

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        mean = self.mean(x)
        std = torch.exp(self.log_std)
        return torch.distributions.Normal(mean, std)

# 策略更新代码片段
old_policy = PolicyNetwork(state_dim, action_dim)
new_policy = PolicyNetwork(state_dim, action_dim)
new_policy.load_state_dict(old_policy.state_dict())

# 计算 KL 散度
kl_div = torch.distributions.kl.kl_divergence(old_policy(states), new_policy(states))
kl_loss = torch.mean(kl_div)

# 加权策略损失
policy_loss = -torch.mean(torch.exp(advantages) * new_policy(states).log_prob(actions))
total_loss = policy_loss + 0.5 * kl_loss  # KL 系数可调

以下是完整的训练循环框架，包含梯度裁剪等稳定化技巧：

# 网络定义
policy_net = PolicyNetwork(state_dim, action_dim)
value_net = ValueNetwork(state_dim)  # 简单 MLP
optimizer = torch.optim.Adam([{'params': policy_net.parameters(), 'lr': 3e-4},
    {'params': value_net.parameters(), 'lr': 1e-3}
])

# 训练循环
for episode in range(1000):
    states, actions, rewards, next_states, dones = env.sample_trajectory()

    # 计算价值估计和优势
    values = value_net(states)
    advantages = compute_advantages(rewards, values, dones)

    # 策略更新
    optimizer.zero_grad()
    policy_loss = -torch.mean(torch.exp(advantages) * policy_net(states).log_prob(actions))
    value_loss = F.mse_loss(values, rewards)  # 简单 MSE 损失

    # 梯度裁剪
    total_loss = policy_loss + value_loss
    total_loss.backward()
    torch.nn.utils.clip_grad_norm_(policy_net.parameters(), 0.5)
    optimizer.step()

通过实验我们发现：