RL-Study/Notebooks/SAC.ipynb

SAC论文复现学习

找了一篇Frank L. Lewis经典的Actor–Critic-Based Optimal Tracking for Partially Unknown Nonlinear Discrete-Time Systems的论文来复现学习，主要是为了理解SAC算法在连续动作空间中的应用，该论文考虑了带约束情况下的最优跟踪问题，使用了Actor-Critic方法来解决。

In [43]:

import torch
import torch.nn as nn
import numpy as np
import math

# 设置随机种子，保证每次运行的结果可复现
torch.manual_seed(42)

# ==========================================
# 超参数与系统参数设置
# ==========================================
gamma = 0.3              # 折扣因子
u_bar = 0.4              # 控制饱和上限 (物理极限)
R = 1.0                  # 控制输入惩罚权重
lr_c = 0.05              # Critic 网络的学习率
lr_a = 0.05              # Actor 网络的学习率

# 状态惩罚矩阵 Q1
# 论文中 Q=20I，由于增广状态 X = [e, r]，且只有误差 e 需要惩罚，所以后两项为0
Q1 = torch.diag(torch.tensor([20.0, 20.0, 0.0, 0.0], dtype=torch.float32))

系统动态方程与参考轨迹发生器

对应的是论文里 Section V 的仿真算例。我们要把离散非线性系统状态方程（公式 70）和正弦命令发生器（公式 71）实现出来 。另外，由于我们后面在 Actor 更新时需要用到输入增益矩阵 $G(X)$，所以这里单独写一个函数提取它。

In [44]:

# ==========================================
# 系统动态与参考轨迹定义
# ==========================================
def system_step(x, u):
    """
    非线性离散时间系统漂移与输入动态
    x: 当前状态向量 [x1, x2]
    u: 控制输入标量
    返回: 下一时刻的状态 [x1_next, x2_next]
    """
    x1_next = -0.8 * x[1]
    x2_next = -0.45 * x[0] - torch.sin(x[1]) + 0.2 * x[1] * u
    return torch.stack([x1_next, x2_next])

def ref_step(r):
    """
    命令发生器动态 (产生振荡参考轨迹)
    r: 当前参考状态向量 [r1, r2]
    返回: 下一时刻的参考状态
    """
    r1_next = -r[0]
    r2_next = -r[1]
    return torch.stack([r1_next, r2_next])

def get_G_matrix(x):
    """
    计算输入增益矩阵 G(X)
    对应于被控系统中控制量 u 前面的状态系数矩阵
    用于 Actor 目标控制量的解析计算
    """
    return torch.tensor([[0.0], [0.2 * x[1]], [0.0], [0.0]], dtype=torch.float32)

构建 Actor 和 Critic 神经网络

论文中采用的是单隐藏层神经网络结构，隐藏层有 5 个神经元，激活函数为 \tanh(.) 。需要特别注意的是，Actor 网络的输出层我们直接加上了 \bar{u} \tanh(.) 的限幅，以此来严格保证输出控制量永远不会超出你设定的物理极限 。

In [45]:

# ==========================================
# 神经网络定义与初始化
# ==========================================
class CriticNet(nn.Module):
    def __init__(self):
        super(CriticNet, self).__init__()
        # Critic网络: 4个输入 -> 5个隐藏层节点(tanh) -> 1个标量输出(线性)
        self.fc1 = nn.Linear(4, 5)
        self.fc2 = nn.Linear(5, 1)
        # 将权重初始化在 -1 到 1 之间
        nn.init.uniform_(self.fc1.weight, -1.0, 1.0)
        nn.init.uniform_(self.fc2.weight, -1.0, 1.0)
        
    def forward(self, X):
        h = torch.tanh(self.fc1(X))
        V = self.fc2(h)
        return V

class ActorNet(nn.Module):
    def __init__(self):
        super(ActorNet, self).__init__()
        # Actor网络: 4个输入 -> 5个隐藏层节点(tanh) -> 1个控制量输出
        self.fc1 = nn.Linear(4, 5)
        self.fc2 = nn.Linear(5, 1)
        nn.init.uniform_(self.fc1.weight, -1.0, 1.0)
        nn.init.uniform_(self.fc2.weight, -1.0, 1.0)
        
    def forward(self, X):
        h = torch.tanh(self.fc1(X))
        # 核心：输出层包含物理饱和约束 u_bar * tanh(.)
        u = u_bar * torch.tanh(self.fc2(h))
        return u

# 实例化网络与对应的随机梯度下降优化器
critic = CriticNet()
actor = ActorNet()
optimizer_c = torch.optim.SGD(critic.parameters(), lr=lr_c)
optimizer_a = torch.optim.SGD(actor.parameters(), lr=lr_a)

定义非二次型性能代价函数

这部分是实现带输入约束控制的核心。通过引入非二次型积分代价 $W(u)$，使得控制算法在优化的同时兼顾输入饱和 。我们用 PyTorch 提供的基础数学算子来实现公式 (20) 的积分解析解。

In [46]:

# ==========================================
# 辅助函数：非二次型积分代价计算
# ==========================================
def non_quadratic_cost(u):
    """
    计算包含输入约束的非二次型积分代价 W(u)
    防止控制量在优化过程中无限制增长
    """
    # 稍微截断控制量，防止在反三角函数中出现数值溢出 (除以0或虚数)
    u_clamped = torch.clamp(u, -u_bar + 1e-5, u_bar - 1e-5)
    ratio = u_clamped / u_bar
    
    # 积分的解析解公式
    integral = (ratio * torch.atanh(ratio)) + 0.5 * torch.log(1 - ratio**2)
    return 2 * (u_bar**2) * R * integral

在线 Actor-Critic 同步学习主循环

这是整个算法的引擎。我们在这里严格贯彻了刚才讨论的时间平移逻辑，利用滞后一步的状态数据来进行贝尔曼误差的计算和权重的反向传播。为了保证系统被充分探索，我们在输出动作上叠加了一个微小的探测噪声。

In [47]:

# ==========================================
# 核心在线学习主循环
# ==========================================
total_steps = 500

# 初始化历史状态序列 (给一个初始偏差，观察跟踪效果)
x_k_minus_1 = torch.tensor([1.0, -1.0], dtype=torch.float32)
r_k_minus_1 = torch.tensor([0.5, -0.5], dtype=torch.float32)

for k in range(1, total_steps + 1):
    # --- 第1部分：前向交互与数据采集 ---
    # 计算 k-1 时刻的误差和增广状态 X(k-1)
    e_k_minus_1 = x_k_minus_1 - r_k_minus_1
    X_k_minus_1 = torch.cat([e_k_minus_1, r_k_minus_1])
    
    # 网络生成策略，并加入探测噪声以满足持续激励(PE)条件
    u_k_minus_1 = actor(X_k_minus_1)
    noise = torch.tensor([0.05 * math.sin(k)], dtype=torch.float32)
    u_applied = torch.clamp(u_k_minus_1 + noise, -u_bar, u_bar) 
    
    # 物理系统演化：得到当前真正的 k 时刻状态
    x_k = system_step(x_k_minus_1, u_applied[0])
    r_k = ref_step(r_k_minus_1)
    e_k = x_k - r_k
    X_k = torch.cat([e_k, r_k])
    
    # --- 第2部分：Critic 网络时移更新 ---
    optimizer_c.zero_grad()
    
    V_k_minus_1 = critic(X_k_minus_1)
    # 截断 k 时刻价值的梯度，把它当作目标常数
    V_k = critic(X_k).detach() 
    
    # 计算历史这一步的真实单步代价
    state_cost = torch.matmul(torch.matmul(X_k_minus_1.T, Q1), X_k_minus_1)
    ctrl_cost = non_quadratic_cost(u_applied[0])
    U_k_minus_1 = state_cost + ctrl_cost

    # 必须将 td_target 彻底切断联系，不让它把梯度传导给 Actor
    td_target = (U_k_minus_1 + gamma * V_k).detach()
    
    # 计算 TD 误差并反向传播更新 Critic
    td_target = U_k_minus_1 + gamma * V_k
    loss_c = 0.5 * (td_target - V_k_minus_1)**2
    loss_c.backward()
    optimizer_c.step()
    
    # --- 第3部分：Actor 网络时移更新 ---
    optimizer_a.zero_grad()
    
    # 重新前向传播以构建计算图
    u_pred = actor(X_k_minus_1)

    # 切断 X_k 与过去系统的联系，仅把它当做一个自变量来求偏导
    X_k_for_grad = X_k.clone().detach().requires_grad_(True)
    V_k_grad = critic(X_k_for_grad)

    # 使用 torch.autograd.grad 纯净地求导，避免对 Critic 权重二次污染
    dV_dX = torch.autograd.grad(outputs=V_k_grad, inputs=X_k_for_grad)[0]
    
    
    # 计算目标控制量 (结合输入增益矩阵 G)
    G_matrix = get_G_matrix(x_k_minus_1)
    inner_term = - (gamma / 2.0) * (1.0 / (u_bar * R)) * torch.matmul(G_matrix.T, dV_dX)
    u_target = u_bar * torch.tanh(inner_term).detach()
    
    # 最小化 Actor 输出与解析目标控制量的误差
    loss_a = 0.5 * (u_target - u_pred)**2
    loss_a.backward()
    optimizer_a.step()
    
    # --- 第4部分：时间轴步进 ---
    x_k_minus_1 = x_k.detach()
    r_k_minus_1 = r_k.detach()
    
    # 每 10 步打印一次学习进度
    if k % 10 == 0:
        print(f"当前步数 {k}: Critic 损失 = {loss_c.item():.4f}, Actor 损失 = {loss_a.item():.4f}")

当前步数 10: Critic 损失 = 664.4067, Actor 损失 = 0.0018
当前步数 20: Critic 损失 = 51.7882, Actor 损失 = 0.0008
当前步数 30: Critic 损失 = 1.2308, Actor 损失 = 0.0004
当前步数 40: Critic 损失 = 0.0579, Actor 损失 = 0.0001
当前步数 50: Critic 损失 = 0.7083, Actor 损失 = 0.0001
当前步数 60: Critic 损失 = 0.7751, Actor 损失 = 0.0000
当前步数 70: Critic 损失 = 0.4341, Actor 损失 = 0.0000
当前步数 80: Critic 损失 = 0.0691, Actor 损失 = 0.0000
当前步数 90: Critic 损失 = 0.0727, Actor 损失 = 0.0000
当前步数 100: Critic 损失 = 0.4509, Actor 损失 = 0.0000
当前步数 110: Critic 损失 = 0.7882, Actor 损失 = 0.0000
当前步数 120: Critic 损失 = 0.6399, Actor 损失 = 0.0000
当前步数 130: Critic 损失 = 0.2181, Actor 损失 = 0.0000
当前步数 140: Critic 损失 = 0.0008, Actor 损失 = 0.0000
当前步数 150: Critic 损失 = 0.2530, Actor 损失 = 0.0000
当前步数 160: Critic 损失 = 0.6776, Actor 损失 = 0.0000
当前步数 170: Critic 损失 = 0.7690, Actor 损失 = 0.0000
当前步数 180: Critic 损失 = 0.4190, Actor 损失 = 0.0000
当前步数 190: Critic 损失 = 0.0433, Actor 损失 = 0.0000
当前步数 200: Critic 损失 = 0.0860, Actor 损失 = 0.0000
当前步数 210: Critic 损失 = 0.4970, Actor 损失 = 0.0000
当前步数 220: Critic 损失 = 0.7905, Actor 损失 = 0.0000
当前步数 230: Critic 损失 = 0.6150, Actor 损失 = 0.0000
当前步数 240: Critic 损失 = 0.1838, Actor 损失 = 0.0000
当前步数 250: Critic 损失 = 0.0039, Actor 损失 = 0.0000
当前步数 260: Critic 损失 = 0.2894, Actor 损失 = 0.0000
当前步数 270: Critic 损失 = 0.7027, Actor 损失 = 0.0000
当前步数 280: Critic 损失 = 0.7510, Actor 损失 = 0.0000
当前步数 290: Critic 损失 = 0.3837, Actor 损失 = 0.0000
当前步数 300: Critic 损失 = 0.0299, Actor 损失 = 0.0000
当前步数 310: Critic 损失 = 0.1111, Actor 损失 = 0.0000
当前步数 320: Critic 损失 = 0.5306, Actor 损失 = 0.0000
当前步数 330: Critic 损失 = 0.7897, Actor 损失 = 0.0000
当前步数 340: Critic 损失 = 0.5876, Actor 损失 = 0.0000
当前步数 350: Critic 损失 = 0.1564, Actor 损失 = 0.0000
当前步数 360: Critic 损失 = 0.0110, Actor 损失 = 0.0000
当前步数 370: Critic 损失 = 0.3220, Actor 损失 = 0.0000
当前步数 380: Critic 损失 = 0.7205, Actor 损失 = 0.0000
当前步数 390: Critic 损失 = 0.7392, Actor 损失 = 0.0000
当前步数 400: Critic 损失 = 0.3486, Actor 损失 = 0.0000
当前步数 410: Critic 损失 = 0.0171, Actor 损失 = 0.0000
当前步数 420: Critic 损失 = 0.1345, Actor 损失 = 0.0000
当前步数 430: Critic 损失 = 0.5621, Actor 损失 = 0.0000
当前步数 440: Critic 损失 = 0.7964, Actor 损失 = 0.0000
当前步数 450: Critic 损失 = 0.5534, Actor 损失 = 0.0000
当前步数 460: Critic 损失 = 0.1282, Actor 损失 = 0.0000
当前步数 470: Critic 损失 = 0.0198, Actor 损失 = 0.0000
当前步数 480: Critic 损失 = 0.3579, Actor 损失 = 0.0000
当前步数 490: Critic 损失 = 0.7433, Actor 损失 = 0.0000
当前步数 500: Critic 损失 = 0.7149, Actor 损失 = 0.0000

In [48]:

import matplotlib.pyplot as plt

# ==========================================
# 算法验证与结果可视化
# ==========================================

# 切换为评估模式 (在此阶段不需要计算梯度)
actor.eval()

# 验证阶段的步数
eval_steps = 500

# 重新初始化状态，用于纯净的闭环测试
x_eval = torch.tensor([1.5, -1.5], dtype=torch.float32)
r_eval = torch.tensor([1.4, -1.4], dtype=torch.float32)

# 用于存放画图数据的列表
x1_history, x2_history = [], []
r1_history, r2_history = [], []
e1_history, e2_history = [], []

# 纯净闭环控制循环 (无梯度更新，无探测噪声)
with torch.no_grad():
    for k in range(eval_steps):
        # 记录当前时刻的数据
        x1_history.append(x_eval[0].item())
        x2_history.append(x_eval[1].item())
        r1_history.append(r_eval[0].item())
        r2_history.append(r_eval[1].item())

        # 记录跟踪误差 e = x - r
        e_eval = x_eval - r_eval
        e1_history.append(e_eval[0].item())
        e2_history.append(e_eval[1].item())

        # 拼接增广状态
        X_eval = torch.cat([e_eval, r_eval])

        # 使用训练好的 Actor 网络计算控制律 (注意：这里去掉了 noise)
        u_eval = actor(X_eval)
        # 加上物理限幅保护
        u_applied = torch.clamp(u_eval, -u_bar, u_bar)

        # 系统状态演化
        x_eval = system_step(x_eval, u_applied[0])
        r_eval = ref_step(r_eval)

# ==========================================
# 使用 Matplotlib 绘制跟踪效果图
# ==========================================
time_steps = range(eval_steps)

# 设置全局字体大小，方便放入大论文
plt.rcParams.update({'font.size': 12})

# 绘制状态 x1 与参考轨迹 r1 的对比图 (对应论文 Fig. 6)
plt.figure(figsize=(10, 5))
plt.plot(time_steps, x1_history, label='x1 (System State)', color='red', linewidth=1.5)
plt.plot(time_steps, r1_history, label='r1 (Reference)', color='green', linestyle='--', linewidth=1.5)
plt.title('Tracking Performance: x1 and r1')
plt.xlabel('Time step')
plt.ylabel('x1 and r1')
plt.legend(loc='upper right')
plt.grid(True)
plt.xlim([0, eval_steps])
plt.show()

# 绘制状态 x2 与参考轨迹 r2 的对比图 (对应论文 Fig. 7)
plt.figure(figsize=(10, 5))
plt.plot(time_steps, x2_history, label='x2 (System State)', color='red', linewidth=1.5)
plt.plot(time_steps, r2_history, label='r2 (Reference)', color='green', linestyle='--', linewidth=1.5)
plt.title('Tracking Performance: x2 and r2')
plt.xlabel('Time step')
plt.ylabel('x2 and r2')
plt.legend(loc='upper right')
plt.grid(True)
plt.xlim([0, eval_steps])
plt.show()

# 新增：绘制误差曲线 e1 和 e2
plt.figure(figsize=(10, 5))
plt.plot(time_steps, e1_history, label='e1 = x1 - r1', color='blue', linewidth=1.5)
plt.plot(time_steps, e2_history, label='e2 = x2 - r2', color='orange', linewidth=1.5)
plt.axhline(0.0, color='black', linestyle='--', linewidth=1.0)
plt.title('Tracking Errors: e1 and e2')
plt.xlabel('Time step')
plt.ylabel('Tracking error')
plt.legend(loc='upper right')
plt.grid(True)
plt.xlim([0, eval_steps])
plt.show()