Adaptive step size for Adam optimizer

Adam stepsize was adjusted based on the target value of the KL divergence

vanilla policy gradient with adaptive stepsize3, After each batch of data, the Adam stepsize is adjusted based on the KL divergence of the original and updated
policy, using a rule similar to the one shown in Section 4. An implementation is available at https://github.com/berkeleydeeprlcourse/homework/tree/master/hw4.

中文精准翻译

带自适应步长的原始策略梯度算法：每收集完一批数据后，依据旧策略与更新后策略之间的 KL 散度，对 Adam 优化器的步长进行自适应调整，所采用的调整规则与第 4 节给出的规则类似。相关实现代码可在以下 GitHub 仓库获取：
https://github.com/berkeleydeeprlcourse/homework/tree/master/hw4

自适应步长策略梯度（Vanilla PG + KL 控制学习率）完整实现
这就是伯克利深度强化学习课程 HW4 的核心算法：
基于策略 KL 散度自动调学习率，策略更新太大就缩小 lr，太小就放大 lr，完美稳定训练。
我直接给你可运行、极简、注释完整的实现。
核心思想（一句话）
用普通策略梯度（REINFORCE）更新策略
计算更新前 vs 更新后策略的 KL 散度
根据 KL 大小自动调整 Adam 学习率：
KL > 目标 → 学习率太大，缩小
KL < 目标 → 学习率太小，放大
KL 正常 → 不动

豆包AI生成下面代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import numpy as np

# ==============================
# 1. 策略网络（最简单的线性/MLP）
# ==============================
class PolicyNet(nn.Module):
    def __init__(self, obs_dim, act_dim):
        super().__init__
        self.fc = nn.Sequential(
            nn.Linear(obs_dim, 64),
            nn.Tanh(),
            nn.Linear(64, act_dim)
        )

    def forward(self, x):
        logits = self.fc(x)
        return logits

# ==============================
# 2. 自适应步长策略梯度（核心）
# ==============================
class AdaptiveStepPG:
    def __init__(
        self,
        obs_dim,
        act_dim,
        lr=1e-2,         # 初始学习率
        target_kl=0.01,  # KL 目标值（最重要超参）
        kl_coeff=1.5     # 学习率缩放系数
    ):
        self.policy = PolicyNet(obs_dim, act_dim)
        self.target_kl = target_kl
        self.kl_coeff = kl_coeff

        # 使用 Adam，**手动控制 lr**
        self.optimizer = optim.Adam(self.policy.parameters(), lr=lr)
        self.lr = lr

    # --------------------------
    # 策略更新 + KL 计算 + 自适应调 lr
    # --------------------------
    def update(self, obs, acts, advs):
        obs = torch.FloatTensor(obs)
        acts = torch.LongTensor(acts)
        advs = torch.FloatTensor(advs)

        # 1. 保存【更新前】的策略分布
        with torch.no_grad():
            old_logits = self.policy(obs)
            old_dist = Categorical(logits=old_logits)

        # 2. 策略梯度损失
        logits = self.policy(obs)
        dist = Categorical(logits=logits)
        log_prob = dist.log_prob(acts)
        loss = -(log_prob * advs).mean()

        # 3. 梯度下降
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # ======================
        # 核心：计算 KL 散度
        # ======================
        with torch.no_grad():
            new_logits = self.policy(obs)
            new_dist = Categorical(logits=new_logits)

            # 计算 KL(old || new)
            kl = torch.mean(
                old_dist.log_prob(acts) - new_dist.log_prob(acts)
            ).item()

        # ======================
        # 自适应调整学习率
        # ======================
        if kl > self.target_kl * 1.5:
            # 更新太大 → 缩小学习率
            self.lr /= self.kl_coeff
            print(f"KL 太大 ({kl:.4f}) → 缩小 LR: {self.lr:.6f}")

        elif kl < self.target_kl / 1.5:
            # 更新太小 → 放大学习率
            self.lr *= self.kl_coeff
            print(f"KL 太小 ({kl:.4f}) → 放大 LR: {self.lr:.6f}")

        # 把新学习率写回优化器
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = self.lr

        return loss.item(), kl

超参数（官方推荐）

target_kl = 0.01 # KL 目标
lr = 0.01 # 初始学习率
kl_coeff = 1.5 # 学习率缩放倍数

训练逻辑（每一批数据调用一次）

# 每一批数据执行一次
loss, kl = agent.update(obs_batch, act_batch, adv_batch)

agent为AdaptiveStepPG类的实例，以上代码为伪代码，示意之用，不能真正运行。