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

1次阅读

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

Advantage-Weighted Regression (AWR) 是一种 off-policy 强化学习算法，相比于 PPO 和 SAC，它在处理连续动作空间时表现出更好的稳定性和样本效率。AWR 的核心思想是通过优势函数（advantage function）来加权策略更新，避免了传统策略梯度方法中的高方差问题。

对比 PPO：AWR 不需要复杂的 clip 机制，减少了超参数调优的负担
对比 SAC：AWR 直接优化策略网络，不需要维护额外的价值函数网络

优先级采样（Prioritized Experience Replay）是 AWR 的关键组件之一。以下是一个基于 PyTorch 的实现：

class PrioritizedReplayBuffer:
    def __init__(self, capacity, alpha=0.6):
        self.capacity = capacity
        self.alpha = alpha  # 优先级调节参数
        self.buffer = []
        self.priorities = np.zeros((capacity,), dtype=np.float32)
        self.pos = 0

    def add(self, state, action, reward, next_state, done):
        max_prio = self.priorities.max() if self.buffer else 1.0

        if len(self.buffer) < self.capacity:
            self.buffer.append((state, action, reward, next_state, done))
        else:
            self.buffer[self.pos] = (state, action, reward, next_state, done)

        self.priorities[self.pos] = max_prio
        self.pos = (self.pos + 1) % self.capacity

    def sample(self, batch_size, beta=0.4):
        if len(self.buffer) == 0:
            return None

        probs = self.priorities[:len(self.buffer)] ** 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()

        return samples, indices, np.array(weights, dtype=np.float32)

优势函数 $A(s,a) = Q(s,a) – V(s)$ 的计算直接影响策略更新的质量。实践中我们使用 GAE（Generalized Advantage Estimation）来平衡偏差和方差：

def compute_advantage(rewards, values, gamma=0.99, lam=0.95):
    """
    rewards: 轨迹上的即时奖励 [T]
    values: 状态价值估计 [T]
    returns: 优势函数估计 [T]
    """
    advantages = np.zeros_like(rewards)
    last_adv = 0

    # 反向计算 GAE
    for t in reversed(range(len(rewards))):
        delta = rewards[t] + gamma * values[t+1] - values[t]
        advantages[t] = last_adv = delta + gamma * lam * last_adv

    # 标准化优势函数
    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.logstd = nn.Parameter(torch.zeros(action_dim))

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return torch.tanh(self.mean(x)), self.logstd.exp()

# KL 约束实现
def kl_constraint(old_probs, new_probs, epsilon=0.02):
    kl = (old_probs * (torch.log(old_probs) - torch.log(new_probs))).sum()
    return torch.max(kl, torch.tensor(epsilon))

以下是整合了所有关键组件的训练循环：

def train_awr(env_name="CartPole-v1", epochs=100, batch_size=64, beta=0.5):
    env = gym.make(env_name)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n if hasattr(env.action_space, 'n') \
                else env.action_space.shape[0]

    # 初始化网络和优化器
    policy = PolicyNetwork(state_dim, action_dim)
    value_net = ValueNetwork(state_dim)
    optimizer = torch.optim.Adam(policy.parameters(), lr=3e-4)

    # 训练循环
    for epoch in range(epochs):
        # 收集轨迹数据
        states, actions, rewards, next_states, dones = collect_trajectories(env, policy)

        # 计算优势函数
        values = value_net(torch.FloatTensor(states))
        advantages = compute_advantage(rewards, values)

        # 策略更新
        for _ in range(5):  # 多次策略更新
            batch_indices = np.random.randint(0, len(states), batch_size)
            batch_states = torch.FloatTensor(states[batch_indices])
            batch_actions = torch.FloatTensor(actions[batch_indices])
            batch_adv = torch.FloatTensor(advantages[batch_indices])

            # 计算新旧策略概率比
            new_log_probs = policy.get_log_prob(batch_states, batch_actions)
            old_log_probs = old_log_probs_buffer[batch_indices]
            ratios = torch.exp(new_log_probs - old_log_probs)

            # 优势加权损失
            loss = -(ratios * batch_adv).mean()

            # KL 约束
            kl = kl_constraint(old_log_probs.exp(), new_log_probs.exp())
            if kl > 0.02:
                optimizer.zero_grad()
                kl.backward()
                torch.nn.utils.clip_grad_norm_(policy.parameters(), 0.5)
                optimizer.step()
                break

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(policy.parameters(), 0.5)
            optimizer.step()

我们对比了不同 β 超参在 CartPole 环境下的表现：