强化学习

import random  
import gymnasium as gym  # 彻底替换为gymnasium，不再用旧gym  
import numpy as np  
import collections  
from tqdm import tqdm  
import torch  
import torch.nn.functional as F  
import matplotlib.pyplot as plt  

# ===== 经验回放池=====  
class ReplayBuffer:  
    ''' 经验回放池 '''    def __init__(self, capacity):  
        self.buffer = collections.deque(maxlen=capacity)  # 队列,先进先出  

    def add(self, state, action, reward, next_state, done):  # 将数据加入buffer  
        self.buffer.append((state, action, reward, next_state, done))  

    def sample(self, batch_size):  # 从buffer中采样数据,数量为batch_size  
        transitions = random.sample(self.buffer, batch_size)  
        state, action, reward, next_state, done = zip(*transitions)  
        return np.array(state), action, reward, np.array(next_state), done  

    def size(self):  # 目前buffer中数据的数量  
        return len(self.buffer)  

# ===== Q网络=====  
class Qnet(torch.nn.Module):  
    ''' 只有一层隐藏层的Q网络 '''    def __init__(self, state_dim, hidden_dim, action_dim):  
        super(Qnet, self).__init__()  
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)  
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)  

    def forward(self, x):  
        x = F.relu(self.fc1(x))  # 隐藏层使用ReLU激活函数  
        return self.fc2(x)  

# ===== DQN算法=====  
class DQN:  
    ''' DQN算法 '''    def __init__(self, state_dim, hidden_dim, action_dim, learning_rate, gamma,  
                 epsilon, target_update, device):  
        self.action_dim = action_dim  
        self.q_net = Qnet(state_dim, hidden_dim,  
                          self.action_dim).to(device)  # Q网络  
        # 目标网络  
        self.target_q_net = Qnet(state_dim, hidden_dim,  
                                 self.action_dim).to(device)  
        # 使用Adam优化器  
        self.optimizer = torch.optim.Adam(self.q_net.parameters(),  
                                          lr=learning_rate)  
        self.gamma = gamma  # 折扣因子  
        self.epsilon = epsilon  # epsilon-贪婪策略  
        self.target_update = target_update  # 目标网络更新频率  
        self.count = 0  # 计数器,记录更新次数  
        self.device = device  

    def take_action(self, state):  # epsilon-贪婪策略采取动作  
        if np.random.random() < self.epsilon:  
            action = np.random.randint(self.action_dim)  
        else:  
            state = torch.tensor([state], dtype=torch.float).to(self.device)  
            action = self.q_net(state).argmax().item()  
        return action  

    def update(self, transition_dict):  
        states = torch.tensor(transition_dict['states'],  
                              dtype=torch.float).to(self.device)  
        actions = torch.tensor(transition_dict['actions']).view(-1, 1).to(  
            self.device)  
        rewards = torch.tensor(transition_dict['rewards'],  
                               dtype=torch.float).view(-1, 1).to(self.device)  
        next_states = torch.tensor(transition_dict['next_states'],  
                                   dtype=torch.float).to(self.device)  
        dones = torch.tensor(transition_dict['dones'],  
                             dtype=torch.float).view(-1, 1).to(self.device)  

        q_values = self.q_net(states).gather(1, actions)  # Q值  
        # 下个状态的最大Q值  
        max_next_q_values = self.target_q_net(next_states).max(1)[0].view(  
            -1, 1)  
        q_targets = rewards + self.gamma * max_next_q_values * (1 - dones  
                                                                )  # TD误差目标  
        dqn_loss = torch.mean(F.mse_loss(q_values, q_targets))  # 均方误差损失函数  
        self.optimizer.zero_grad()  # PyTorch中默认梯度会累积,这里需要显式将梯度置为0  
        dqn_loss.backward()  # 反向传播更新参数  
        self.optimizer.step()  

        if self.count % self.target_update == 0:  
            self.target_q_net.load_state_dict(  
                self.q_net.state_dict())  # 更新目标网络  
        self.count += 1  

# ===== 主程序（小车悬臂的问题）=====  
if __name__ == "__main__":  
    # 1. 配置参数  
    lr = 2e-3               # 学习率  
    num_episodes = 500      # 训练总回合数  
    hidden_dim = 128        # Q网络隐藏层维度  
    gamma = 0.98            # 折扣因子  
    epsilon = 0.01          # epsilon-贪婪策略  
    target_update = 10      # 目标网络更新频率  
    buffer_size = 10000     # 经验回放池容量  
    batch_size = 64         # 批次大小  
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")  

    # 2. 创建环境（修正1：用CartPole-v1，移除seed()，改用reset设置种子）  
    env = gym.make("CartPole-v1", render_mode="human")  # render_mode="human" 方便测试时可视化  
    state_dim = env.observation_space.shape[0]  
    action_dim = env.action_space.n  

    # 3. 初始化组件  
    replay_buffer = ReplayBuffer(buffer_size)  
    agent = DQN(state_dim, hidden_dim, action_dim, lr, gamma, epsilon, target_update, device)  

    # 4. 训练循环  
    return_list = []  
    for i in range(10):  
        with tqdm(total=int(num_episodes/10), desc='Iteration %d' % i) as pbar:  
            for i_episode in range(int(num_episodes/10)):  
                episode_return = 0  
                # 修正2：用reset(seed=0)设置随机种子，替代旧的env.seed(0)  
                state, _ = env.reset(seed=0)  # Gymnasium的reset返回(state, info)，并支持seed参数  
                done = False  
                while not done:  
                    action = agent.take_action(state)  
                    # 修正3：Gymnasium的step返回(next_state, reward, terminated, truncated, info)  
                    next_state, reward, terminated, truncated, _ = env.step(action)  
                    # 合并终止条件（terminated是任务完成，truncated是步数超限）  
                    done = terminated or truncated  
                    replay_buffer.add(state, action, reward, next_state, done)  
                    state = next_state  
                    episode_return += reward  

                    if replay_buffer.size() > batch_size:  
                        b_s, b_a, b_r, b_ns, b_d = replay_buffer.sample(batch_size)  
                        transition_dict = {  
                            'states': b_s,  
                            'actions': b_a,  
                            'next_states': b_ns,  
                            'rewards': b_r,  
                            'dones': b_d  
                        }  
                        agent.update(transition_dict)  
                return_list.append(episode_return)  

                if (i_episode+1) % 10 == 0:  
                    pbar.set_postfix({  
                        'episode':  
                        '%d' % (num_episodes/10 * i + i_episode + 1),  
                        'return':  
                        '%.3f' % np.mean(return_list[-10:])  
                    })  
                pbar.update(1)  

    # 5. 可视化奖励曲线  
    episodes_list = list(range(len(return_list)))  
    plt.plot(episodes_list, return_list)  
    plt.xlabel('Episodes')  
    plt.ylabel('Returns')  
    plt.title('DQN on CartPole-v1')  
    plt.show()  

    # 6. 测试智能体（可选）  
    test_episodes = 5  
    for _ in range(test_episodes):  
        state, _ = env.reset(seed=0)  
        done = False  
        while not done:  
            action = agent.take_action(state)  
            next_state, reward, terminated, truncated, _ = env.step(action)  
            done = terminated or truncated  
            state = next_state  
    env.close()

class ConvolutionalQnet(torch.nn.Module):
    ''' 加入卷积层的Q网络 '''
    def __init__(self, action_dim, in_channels=4):
        super(ConvolutionalQnet, self).__init__()
        self.conv1 = torch.nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
        self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=4, stride=2)
        self.conv3 = torch.nn.Conv2d(64, 64, kernel_size=3, stride=1)
        self.fc4 = torch.nn.Linear(7 * 7 * 64, 512)
        self.head = torch.nn.Linear(512, action_dim)

    def forward(self, x):
        x = x / 255
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.fc4(x))
        return self.head(x)

import gymnasium as gym  
import torch  
import torch.nn.functional as F  
import numpy as np  
import matplotlib.pyplot as plt  
from tqdm import tqdm  
import random  # 新增：用于设置全局随机种子  


# ===== 策略网络（不变）=====  
class PolicyNet(torch.nn.Module):  
    def __init__(self, state_dim, hidden_dim, action_dim):  
        super(PolicyNet, self).__init__()  
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)  
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)  

    def forward(self, x):  
        x = F.relu(self.fc1(x))  
        return F.softmax(self.fc2(x), dim=1)  


# ===== REINFORCE算法（不变）=====  
class REINFORCE:  
    def __init__(self, state_dim, hidden_dim, action_dim, learning_rate, gamma,  
                 device):  
        self.policy_net = PolicyNet(state_dim, hidden_dim,  
                                    action_dim).to(device)  
        self.optimizer = torch.optim.Adam(self.policy_net.parameters(),  
                                          lr=learning_rate)  # 使用Adam优化器  
        self.gamma = gamma  # 折扣因子  
        self.device = device  

    def take_action(self, state):  # 根据动作概率分布随机采样  
        state = torch.tensor([state], dtype=torch.float).to(self.device)  
        probs = self.policy_net(state)  
        action_dist = torch.distributions.Categorical(probs)  
        action = action_dist.sample()  
        return action.item()  

    def update(self, transition_dict):  
        reward_list = transition_dict['rewards']  
        state_list = transition_dict['states']  
        action_list = transition_dict['actions']  

        G = 0  
        self.optimizer.zero_grad()  
        for i in reversed(range(len(reward_list))):  # 从最后一步算起  
            reward = reward_list[i]  
            state = torch.tensor([state_list[i]],  
                                 dtype=torch.float).to(self.device)  
            action = torch.tensor([action_list[i]]).view(-1, 1).to(self.device)  
            log_prob = torch.log(self.policy_net(state).gather(1, action))  
            G = self.gamma * G + reward  
            loss = -log_prob * G  # 每一步的损失函数  
            loss.backward()  # 反向传播计算梯度  
        self.optimizer.step()  # 梯度下降  


# ===== 主程序（核心修正部分）=====  
if __name__ == "__main__":  
    # 1. 配置参数（不变）  
    learning_rate = 1e-3  
    num_episodes = 1000  
    hidden_dim = 128  
    gamma = 0.98  
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")  


    # 2. 设置全局随机种子（新增：保证实验可复现，替代原env.seed(0)）  
    def set_seed(seed=0):  
        random.seed(seed)  
        np.random.seed(seed)  
        torch.manual_seed(seed)  
        if torch.cuda.is_available():  
            torch.cuda.manual_seed(seed)  
            torch.cuda.manual_seed_all(seed)  


    set_seed(0)  

    # 3. 创建环境（修正1：升级为CartPole-v1，避免过时警告）  
    env_name = "CartPole-v1"  
    env = gym.make(env_name)  

    # 4. 移除env.seed(0)（核心修正2：该方法已废弃）  
    # env.seed(0)  # 删掉这行！  

    # 5. 获取环境维度（不变）  
    state_dim = env.observation_space.shape[0]  
    action_dim = env.action_space.n  
    agent = REINFORCE(state_dim, hidden_dim, action_dim, learning_rate, gamma, device)  

    # 6. 训练循环（核心修正3：适配reset/step返回值）  
    return_list = []  
    for i in range(10):  
        with tqdm(total=int(num_episodes / 10), desc='Iteration %d' % i) as pbar:  
            for i_episode in range(int(num_episodes / 10)):  
                episode_return = 0  
                transition_dict = {  
                    'states': [],  
                    'actions': [],  
                    'next_states': [],  
                    'rewards': [],  
                    'dones': []  
                }  
                # 修正4：Gymnasium的reset返回(state, info)，并设置seed保证可复现  
                state, info = env.reset(seed=0)  
                done = False  
                while not done:  
                    action = agent.take_action(state)  
                    # 修正5：Gymnasium的step返回5个值（next_state, reward, terminated, truncated, info）  
                    next_state, reward, terminated, truncated, info = env.step(action)  
                    # 修正6：合并终止条件（terminated=任务完成，truncated=步数超限）  
                    done = terminated or truncated  

                    # 以下逻辑不变  
                    transition_dict['states'].append(state)  
                    transition_dict['actions'].append(action)  
                    transition_dict['next_states'].append(next_state)  
                    transition_dict['rewards'].append(reward)  
                    transition_dict['dones'].append(done)  
                    state = next_state  
                    episode_return += reward  
                return_list.append(episode_return)  
                agent.update(transition_dict)  
                if (i_episode + 1) % 10 == 0:  
                    pbar.set_postfix({  
                        'episode':  
                            '%d' % (num_episodes / 10 * i + i_episode + 1),  
                        'return':  
                            '%.3f' % np.mean(return_list[-10:])  
                    })  
                pbar.update(1)  

    # 新增：可视化训练结果（可选，方便看效果）  
    episodes_list = list(range(len(return_list)))  
    plt.plot(episodes_list, return_list)  
    plt.xlabel('Episodes')  
    plt.ylabel('Returns')  
    plt.title('REINFORCE on CartPole-v1')  
    plt.show()  

    # 关闭环境  
    env.close()

import gymnasium as gym  
import torch  
import torch.nn.functional as F  
import numpy as np  
import matplotlib.pyplot as plt  
from tqdm import tqdm  
import random  # 新增：设置全局随机种子  

# ===== 1. 补全缺失的ValueNet类（核心！原代码未定义）=====  
class ValueNet(torch.nn.Module):  
    def __init__(self, state_dim, hidden_dim):  
        super(ValueNet, self).__init__()  
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)  
        self.fc2 = torch.nn.Linear(hidden_dim, 1)  # 输出状态价值（标量）  

    def forward(self, x):  
        x = F.relu(self.fc1(x))  
        return self.fc2(x)  # 输出V(s)  

# ===== 2. 策略网络（不变）=====  
class PolicyNet(torch.nn.Module):  
    def __init__(self, state_dim, hidden_dim, action_dim):  
        super(PolicyNet, self).__init__()  
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)  
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)  

    def forward(self, x):  
        x = F.relu(self.fc1(x))  
        return F.softmax(self.fc2(x), dim=1)  # 输出动作概率分布  

# ===== 3. Actor-Critic算法（仅适配设备，逻辑不变）=====  
class ActorCritic:  
    def __init__(self, state_dim, hidden_dim, action_dim, actor_lr, critic_lr,  
                 gamma, device):  
        # Actor：策略网络（选动作）  
        self.actor = PolicyNet(state_dim, hidden_dim, action_dim).to(device)  
        # Critic：价值网络（评估状态价值）  
        self.critic = ValueNet(state_dim, hidden_dim).to(device)  
        # 优化器（分开优化Actor和Critic）  
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)  
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=critic_lr)  
        self.gamma = gamma  # 折扣因子  
        self.device = device  

    def take_action(self, state):  # 根据策略网络采样动作  
        state = torch.tensor([state], dtype=torch.float).to(self.device)  
        probs = self.actor(state)  
        action_dist = torch.distributions.Categorical(probs)  
        action = action_dist.sample()  
        return action.item()  

    def update(self, transition_dict):  
        # 解析经验数据  
        states = torch.tensor(transition_dict['states'], dtype=torch.float).to(self.device)  
        actions = torch.tensor(transition_dict['actions']).view(-1, 1).to(self.device)  
        rewards = torch.tensor(transition_dict['rewards'], dtype=torch.float).view(-1, 1).to(self.device)  
        next_states = torch.tensor(transition_dict['next_states'], dtype=torch.float).to(self.device)  
        dones = torch.tensor(transition_dict['dones'], dtype=torch.float).view(-1, 1).to(self.device)  

        # Critic更新：时序差分（TD）误差优化价值网络  
        td_target = rewards + self.gamma * self.critic(next_states) * (1 - dones)  # TD目标  
        td_delta = td_target - self.critic(states)  # TD误差（优势函数近似）  
        critic_loss = torch.mean(F.mse_loss(self.critic(states), td_target.detach()))  # 均方误差损失  

        # Actor更新：用TD误差指导策略梯度（降低方差）  
        log_probs = torch.log(self.actor(states).gather(1, actions))  
        actor_loss = torch.mean(-log_probs * td_delta.detach())  # 策略梯度损失  

        # 反向传播+参数更新  
        self.actor_optimizer.zero_grad()  
        self.critic_optimizer.zero_grad()  
        actor_loss.backward()  
        critic_loss.backward()  
        self.actor_optimizer.step()  
        self.critic_optimizer.step()  

# ===== 4. 替代rl_utils.train_on_policy_agent（自定义训练循环，避免依赖缺失）=====  
def train_on_policy_agent(env, agent, num_episodes):  
    return_list = []  # 记录每回合奖励  
    for i in range(10):  # 分10个迭代段，用tqdm显示进度  
        with tqdm(total=int(num_episodes/10), desc=f'Iteration {i}') as pbar:  
            for i_episode in range(int(num_episodes/10)):  
                episode_return = 0  
                transition_dict = {'states': [], 'actions': [], 'next_states': [], 'rewards': [], 'dones': []}  
                # 重置环境（适配Gymnasium：返回(state, info)，设置seed保证可复现）  
                state, info = env.reset(seed=0)  
                done = False  
                while not done:  
                    action = agent.take_action(state)  
                    # 执行动作（适配Gymnasium：step返回5个值）  
                    next_state, reward, terminated, truncated, info = env.step(action)  
                    done = terminated or truncated  # 合并终止条件  
                    # 存储经验  
                    transition_dict['states'].append(state)  
                    transition_dict['actions'].append(action)  
                    transition_dict['next_states'].append(next_state)  
                    transition_dict['rewards'].append(reward)  
                    transition_dict['dones'].append(done)  
                    # 更新状态和奖励  
                    state = next_state  
                    episode_return += reward  
                # 记录奖励+更新策略（On-Policy：每回合结束后更新）  
                return_list.append(episode_return)  
                agent.update(transition_dict)  
                # 进度条更新  
                if (i_episode+1) % 10 == 0:  
                    pbar.set_postfix({'episode': f'{num_episodes/10*i + i_episode+1}', 'return': f'{np.mean(return_list[-10:]):.3f}'})  
                pbar.update(1)  
    return return_list  

# ===== 5. 主程序（核心修正：适配Gymnasium+补全配置）=====  
if __name__ == "__main__":  
    # 超参数配置  
    actor_lr = 1e-3  
    critic_lr = 1e-2  
    num_episodes = 1000  
    hidden_dim = 128  
    gamma = 0.98  
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")  

    # 设置全局随机种子（替代env.seed(0)，保证实验可复现）  
    def set_seed(seed=0):  
        random.seed(seed)  
        np.random.seed(seed)  
        torch.manual_seed(seed)  
        if torch.cuda.is_available():  
            torch.cuda.manual_seed(seed)  
            torch.cuda.manual_seed_all(seed)  
    set_seed(0)  

    # 创建环境（升级为CartPole-v1，避免过时警告）  
    env_name = 'CartPole-v1'  
    env = gym.make(env_name)  
    state_dim = env.observation_space.shape[0]  
    action_dim = env.action_space.n  

    # 初始化智能体+开始训练  
    agent = ActorCritic(state_dim, hidden_dim, action_dim, actor_lr, critic_lr, gamma, device)  
    return_list = train_on_policy_agent(env, agent, num_episodes)  

    # 可视化训练结果  
    episodes_list = list(range(len(return_list)))  
    plt.plot(episodes_list, return_list)  
    plt.xlabel('Episodes')  
    plt.ylabel('Returns')  
    plt.title('Actor-Critic on CartPole-v1')  
    plt.show()  

    # 关闭环境  
    env.close()

import torch  
import numpy as np  
import gymnasium as gym  
import matplotlib.pyplot as plt  
import torch.nn.functional as F  
import copy  
from tqdm import tqdm  

# 仅补全缺失的ValueNet（原代码调用但未定义）  
class ValueNet(torch.nn.Module):  
    def __init__(self, state_dim, hidden_dim):  
        super(ValueNet, self).__init__()  
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)  
        self.fc2 = torch.nn.Linear(hidden_dim, 1)  

    def forward(self, x):  
        x = F.relu(self.fc1(x))  
        return self.fc2(x)  

# 仅补全缺失的compute_advantage（原代码调用但未定义）  
def compute_advantage(gamma, lmbda, td_delta):  
    td_delta = td_delta.detach().numpy()  
    advantage_list = []  
    advantage = 0.0  
    for delta in reversed(td_delta):  
        advantage = gamma * lmbda * advantage + delta  
        advantage_list.append(advantage)  
    advantage_list.reverse()  
    return torch.tensor(advantage_list, dtype=torch.float)  

# 保留你原代码的PolicyNetContinuous  
class PolicyNetContinuous(torch.nn.Module):  
    def __init__(self, state_dim, hidden_dim, action_dim):  
        super(PolicyNetContinuous, self).__init__()  
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)  
        self.fc_mu = torch.nn.Linear(hidden_dim, action_dim)  
        self.fc_std = torch.nn.Linear(hidden_dim, action_dim)  

    def forward(self, x):  
        x = F.relu(self.fc1(x))  
        mu = 2.0 * torch.tanh(self.fc_mu(x))  
        std = F.softplus(self.fc_std(x))  
        return mu, std  

# 保留你原代码的TRPOContinuous  
class TRPOContinuous:  
    def __init__(self, hidden_dim, state_space, action_space, lmbda,  
                 kl_constraint, alpha, critic_lr, gamma, device):  
        state_dim = state_space.shape[0]  
        action_dim = action_space.shape[0]  
        self.actor = PolicyNetContinuous(state_dim, hidden_dim,  
                                         action_dim).to(device)  
        self.critic = ValueNet(state_dim, hidden_dim).to(device)  
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),  
                                                 lr=critic_lr)  
        self.gamma = gamma  
        self.lmbda = lmbda  
        self.kl_constraint = kl_constraint  
        self.alpha = alpha  
        self.device = device  

    def take_action(self, state):  
        state = torch.tensor([state], dtype=torch.float).to(self.device)  
        mu, std = self.actor(state)  
        action_dist = torch.distributions.Normal(mu, std)  
        action = action_dist.sample()  
        return [action.item()]  

    def hessian_matrix_vector_product(self,  
                                      states,  
                                      old_action_dists,  
                                      vector,  
                                      damping=0.1):  
        mu, std = self.actor(states)  
        new_action_dists = torch.distributions.Normal(mu, std)  
        kl = torch.mean(  
            torch.distributions.kl.kl_divergence(old_action_dists,  
                                                 new_action_dists))  
        kl_grad = torch.autograd.grad(kl,  
                                      self.actor.parameters(),  
                                      create_graph=True)  
        kl_grad_vector = torch.cat([grad.view(-1) for grad in kl_grad])  
        kl_grad_vector_product = torch.dot(kl_grad_vector, vector)  
        grad2 = torch.autograd.grad(kl_grad_vector_product,  
                                    self.actor.parameters())  
        grad2_vector = torch.cat(  
            [grad.contiguous().view(-1) for grad in grad2])  
        return grad2_vector + damping * vector  

    def conjugate_gradient(self, grad, states, old_action_dists):  
        x = torch.zeros_like(grad)  
        r = grad.clone()  
        p = grad.clone()  
        rdotr = torch.dot(r, r)  
        for i in range(10):  
            Hp = self.hessian_matrix_vector_product(states, old_action_dists,  
                                                    p)  
            alpha = rdotr / torch.dot(p, Hp)  
            x += alpha * p  
            r -= alpha * Hp  
            new_rdotr = torch.dot(r, r)  
            if new_rdotr < 1e-10:  
                break  
            beta = new_rdotr / rdotr  
            p = r + beta * p  
            rdotr = new_rdotr  
        return x  

    def compute_surrogate_obj(self, states, actions, advantage, old_log_probs,  
                              actor):  
        mu, std = actor(states)  
        action_dists = torch.distributions.Normal(mu, std)  
        log_probs = action_dists.log_prob(actions)  
        ratio = torch.exp(log_probs - old_log_probs)  
        return torch.mean(ratio * advantage)  

    def line_search(self, states, actions, advantage, old_log_probs,  
                    old_action_dists, max_vec):  
        old_para = torch.nn.utils.convert_parameters.parameters_to_vector(  
            self.actor.parameters())  
        old_obj = self.compute_surrogate_obj(states, actions, advantage,  
                                             old_log_probs, self.actor)  
        for i in range(15):  
            coef = self.alpha**i  
            new_para = old_para + coef * max_vec  
            new_actor = copy.deepcopy(self.actor)  
            torch.nn.utils.convert_parameters.vector_to_parameters(  
                new_para, new_actor.parameters())  
            mu, std = new_actor(states)  
            new_action_dists = torch.distributions.Normal(mu, std)  
            kl_div = torch.mean(  
                torch.distributions.kl.kl_divergence(old_action_dists,  
                                                     new_action_dists))  
            new_obj = self.compute_surrogate_obj(states, actions, advantage,  
                                                 old_log_probs, new_actor)  
            if new_obj > old_obj and kl_div < self.kl_constraint:  
                return new_para  
        return old_para  

    def policy_learn(self, states, actions, old_action_dists, old_log_probs,  
                     advantage):  
        surrogate_obj = self.compute_surrogate_obj(states, actions, advantage,  
                                                   old_log_probs, self.actor)  
        grads = torch.autograd.grad(surrogate_obj, self.actor.parameters())  
        obj_grad = torch.cat([grad.view(-1) for grad in grads]).detach()  
        descent_direction = self.conjugate_gradient(obj_grad, states,  
                                                    old_action_dists)  
        Hd = self.hessian_matrix_vector_product(states, old_action_dists,  
                                                descent_direction)  
        max_coef = torch.sqrt(2 * self.kl_constraint /  
                              (torch.dot(descent_direction, Hd) + 1e-8))  
        new_para = self.line_search(states, actions, advantage, old_log_probs,  
                                    old_action_dists,  
                                    descent_direction * max_coef)  
        torch.nn.utils.convert_parameters.vector_to_parameters(  
            new_para, self.actor.parameters())  

    def update(self, transition_dict):  
        states = torch.tensor(transition_dict['states'],  
                              dtype=torch.float).to(self.device)  
        actions = torch.tensor(transition_dict['actions'],  
                               dtype=torch.float).view(-1, 1).to(self.device)  
        rewards = torch.tensor(transition_dict['rewards'],  
                               dtype=torch.float).view(-1, 1).to(self.device)  
        next_states = torch.tensor(transition_dict['next_states'],  
                                   dtype=torch.float).to(self.device)  
        dones = torch.tensor(transition_dict['dones'],  
                             dtype=torch.float).view(-1, 1).to(self.device)  
        rewards = (rewards + 8.0) / 8.0  
        td_target = rewards + self.gamma * self.critic(next_states) * (1 -  
                                                                       dones)  
        td_delta = td_target - self.critic(states)  
        advantage = compute_advantage(self.gamma, self.lmbda,  
                                      td_delta.cpu()).to(self.device)  
        mu, std = self.actor(states)  
        old_action_dists = torch.distributions.Normal(mu.detach(),  
                                                      std.detach())  
        old_log_probs = old_action_dists.log_prob(actions)  
        critic_loss = torch.mean(  
            F.mse_loss(self.critic(states), td_target.detach()))  
        self.critic_optimizer.zero_grad()  
        critic_loss.backward()  
        self.critic_optimizer.step()  
        self.policy_learn(states, actions, old_action_dists, old_log_probs,  
                          advantage)  

# 仅替换rl_utils.train_on_policy_agent（极简版，适配Gymnasium）  
def train_on_policy_agent(env, agent, num_episodes):  
    return_list = []  
    for i in range(10):  
        with tqdm(total=int(num_episodes/10), desc=f'Iteration {i}') as pbar:  
            for i_episode in range(int(num_episodes/10)):  
                episode_return = 0  
                transition_dict = {'states': [], 'actions': [], 'next_states': [], 'rewards': [], 'dones': []}  
                # 仅修改reset：适配Gymnasium返回值+设置seed  
                state, info = env.reset(seed=0)  
                done = False  
                while not done:  
                    action = agent.take_action(state)  
                    # 仅修改step：适配Gymnasium返回值  
                    next_state, reward, terminated, truncated, info = env.step(action)  
                    done = terminated or truncated  
                    transition_dict['states'].append(state)  
                    transition_dict['actions'].append(action)  
                    transition_dict['next_states'].append(next_state)  
                    transition_dict['rewards'].append(reward)  
                    transition_dict['dones'].append(done)  
                    state = next_state  
                    episode_return += reward  
                return_list.append(episode_return)  
                agent.update(transition_dict)  
                if (i_episode + 1) % 10 == 0:  
                    pbar.set_postfix({'episode': f'{num_episodes/10*i + i_episode+1}', 'return': f'{np.mean(return_list[-10:]):.3f}'})  
                pbar.update(1)  
    return return_list  

# 仅替换rl_utils.moving_average（极简版）  
def moving_average(x, window_size):  
    return np.convolve(x, np.ones(window_size)/window_size, mode='valid')  

# 保留你原代码的超参数和主逻辑，仅删除env.seed(0)  
num_episodes = 2000  
hidden_dim = 128  
gamma = 0.9  
lmbda = 0.9  
critic_lr = 1e-2  
kl_constraint = 0.00005  
alpha = 0.5  
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")  

env_name = 'Pendulum-v1'  
env = gym.make(env_name)  
# 仅删除：env.seed(0)  
torch.manual_seed(0)  
agent = TRPOContinuous(hidden_dim, env.observation_space, env.action_space,  
                       lmbda, kl_constraint, alpha, critic_lr, gamma, device)  
return_list = train_on_policy_agent(env, agent, num_episodes)  

episodes_list = list(range(len(return_list)))  
plt.plot(episodes_list, return_list)  
plt.xlabel('Episodes')  
plt.ylabel('Returns')  
plt.title('TRPO on {}'.format(env_name))  
plt.show()  

# 适配moving_average输出长度（仅小修改）  
mv_return = moving_average(return_list, 9)  
plt.plot(episodes_list[4:-4], mv_return)  # 仅调整x轴长度匹配  
plt.xlabel('Episodes')  
plt.ylabel('Returns')  
plt.title('TRPO on {}'.format(env_name))  
plt.show()