update

examples/playground/chat.py

@@ -17,7 +17,7 @@ def get_args():
     parser.add_argument(
         "--pretrained_model_name_or_path",
         # default="jingyaogong/MiniMind2",
-        default=(project_path / "pretrained_models/MiniMind2"),
+        default=(project_path / "pretrained_models/jingyaogong/MiniMind2"),
         type=str
     )
 

examples/playground/generation.py

@@ -0,0 +1,78 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/jingyaogong/minimind/blob/master/eval_llm.py
+"""
+import argparse
+import time
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+
+from project_settings import project_path
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        # default=(project_path / "trained_models/gpt2-sst2-generation"),
+        default=(project_path / "trained_models/gpt2-sst2-generation-20260213-2048"),
+        type=str
+    )
+    parser.add_argument(
+        "--max_new_tokens",
+        default=1024, # 8192, 128
+        type=int, help="最大生成长度（注意：并非模型实际长文本能力）"
+    )
+    parser.add_argument("--top_p", default=0.85, type=float, help="nucleus采样阈值（0-1）")
+    parser.add_argument("--temperature", default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        # device = "mps"
+        device = "cpu"
+    else:
+        device = "cpu"
+    print(f"device: {device}")
+
+    tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name_or_path)
+    model = AutoModelForCausalLM.from_pretrained(args.pretrained_model_name_or_path)
+    model = model.eval().to(device)
+
+    tokenized = tokenizer(
+        # "this",
+        # "this is ",
+        # "who needs mind-bending",
+        "eldom has a movie",
+        # "thanks to scott 's charismatic",
+        return_tensors="pt"
+    )
+
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    generated_ids = model.generate(
+        inputs=tokenized["input_ids"], attention_mask=tokenized["attention_mask"],
+        max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
+        pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
+        top_p=args.top_p, temperature=args.temperature, repetition_penalty=3.0,
+        early_stopping=True,
+    )
+    # response = tokenizer.decode(generated_ids[0][len(tokenized["input_ids"][0]):], skip_special_tokens=True)
+    response = tokenizer.decode(generated_ids[0], skip_special_tokens=False)
+    print(response)
+    print(generated_ids)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/dpo/ultrachat/step_1_prepare_data.py

@@ -0,0 +1,59 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+或使用命令行
+pip install modelscope
+modelscope download \
+--model 'qgyd2021/Qwen3-8B-sft-deepspeed' \
+--local_dir '/root/autodl-tmp/trained_models/Qwen3-8B-sft-deepspeed'
+
+"""
+import argparse
+import os
+from pathlib import Path
+import platform
+
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from modelscope import snapshot_download
+# from huggingface_hub import snapshot_download
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo_id", default="openai-community/gpt2", type=str)
+    parser.add_argument(
+        "--local_dir",
+        default=(temp_directory / "../trained_models/openai-community/gpt2").as_posix(),
+        type=str
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    #modelscope
+    snapshot_download(
+        model_id=args.repo_id,
+        local_dir=args.local_dir,
+    )
+    #huggingface_hub
+    snapshot_download(
+        repo_type="model",
+        repo_id=args.repo_id,
+        local_dir=args.local_dir,
+    )
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/dpo/ultrachat/step_2_train_sft_model.py

@@ -0,0 +1,132 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import os
+from pathlib import Path
+import platform
+
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from datasets import load_dataset
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
+from trl import SFTTrainer, SFTConfig
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_name",
+        default=(project_path / "pretrained_models/Qwen/Qwen2.5-0.5B").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--dataset_path",
+        default="HuggingFaceH4/ultrachat_200k",
+        type=str
+    ),
+    parser.add_argument("--dataset_name", default=None, type=str),
+    parser.add_argument("--dataset_split", default=None, type=str),
+    parser.add_argument(
+        "--dataset_cache_dir",
+        default=(temp_directory / "hub_datasets").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--model_cache_dir",
+        default=(temp_directory / "hub_models").as_posix(),
+        type=str
+    ),
+    parser.add_argument("--dataset_streaming", default="false", type=str),
+    parser.add_argument("--valid_dataset_size", default=1000, type=int),
+    parser.add_argument("--shuffle_buffer_size", default=5000, type=int),
+
+    parser.add_argument("--max_seq_length", default=2048, type=int)
+
+    parser.add_argument(
+        "--output_model_dir",
+        default=(project_path / "trained_models/qwen2_5-0_5B-ultrachat-sft").as_posix(),
+        type=str
+    ),
+
+    parser.add_argument(
+        "--num_workers",
+        default=None if platform.system() in ("Windows", "Darwin") else os.cpu_count() // 2,
+        type=int
+    ),
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = AutoModelForCausalLM.from_pretrained(args.model_name)
+    model = model.to(args.device)
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    train_dataset = load_dataset(
+        path=args.dataset_path,
+        name=args.dataset_name,
+        split="train_sft",
+        cache_dir=args.dataset_cache_dir,
+        # num_proc=args.num_workers if not args.dataset_streaming else None,
+        streaming=True if args.dataset_streaming in ("true",) else False,
+    )
+
+    sft_config = SFTConfig(
+        output_dir=args.output_model_dir,
+        num_train_epochs=1,
+        per_device_train_batch_size=4,
+        gradient_accumulation_steps=4,
+        save_strategy="steps",
+        save_steps=500,
+        save_total_limit=2,
+        logging_steps=100,
+        learning_rate=2e-5,
+        warmup_ratio=0.03,
+        lr_scheduler_type="cosine",
+        bf16=torch.cuda.is_available(),
+        tf32=torch.cuda.is_available(),
+        gradient_checkpointing=True,
+        optim="adamw_torch",
+        remove_unused_columns=False,
+        report_to="none",
+        dataloader_num_workers=args.num_workers or 0,
+        ddp_find_unused_parameters=False if torch.cuda.device_count() > 1 else None,
+
+        # SFT specific parameters
+        max_length=args.max_seq_length,
+        dataset_text_field=None,
+        dataset_kwargs={
+            "add_special_tokens": True,
+            # "split": "train",
+        },
+    )
+
+    # 创建 trainer
+    trainer = SFTTrainer(
+        model=model,
+        args=sft_config,
+        train_dataset=train_dataset,
+        processing_class=tokenizer,
+    )
+
+    trainer.train()
+    trainer.save_model()
+    return
+
+
+if __name__ == "__main__":
+    main()
+

examples/tutorials/mix_lora_unsloth/step_2_train_model.py

@@ -146,6 +146,8 @@ def main():
             # max_steps = 30,
             learning_rate=2e-5,  # Reduce to 2e-5 for long training runs
             logging_steps=1,
+            save_steps=100,  # 每500步保存一次检查点
+            save_total_limit=2,  # 最多只保留2个检查点，旧的自动清理
             optim="adamw_8bit",
             weight_decay=0.01,
             lr_scheduler_type="linear",

examples/tutorials/rl/cart_pole/requirements.txt

@@ -0,0 +1,3 @@
+gymnasium
+matplotlib
+pygame

examples/tutorials/rl/cart_pole/step_2_actor_critic.py

@@ -0,0 +1,716 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+相比于 REINFORCE 方法，演员-评论家（Actor-Critic）方法不需要等到一局游戏结束就可以触发优化迭代。
+
+
+
+A2C 计算价值优势，主要是对单次的优势进行了移动平均。
+
+由于函数的优化步长受到优势的直接影响。当优步长过大时很容易直接跨过，导致优化失败。
+虽然A2C 已经对历史优势进行了移动平均，但问题仍然存在。
+尤其是当训练的早期价值函数还没有获得较好的训练，这种问题尤其容易出现。
+因此需要对优化步长进行截断限制。
+
+"""
+import gymnasium as gym
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.distributions import Categorical
+import torch.nn.functional as F
+
+
+# ============== 1. 基础Actor-Critic ==============
+class ActorCritic(nn.Module):
+    """共享网络的Actor-Critic"""
+
+    def __init__(self, state_dim, action_dim, hidden_dim=256):
+        super(ActorCritic, self).__init__()
+
+        # 共享特征提取层
+        self.feature_layer = nn.Sequential(
+            nn.Linear(state_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU()
+        )
+
+        # Actor: 策略网络
+        self.actor = nn.Linear(hidden_dim, action_dim)
+
+        # Critic: 价值网络
+        self.critic = nn.Linear(hidden_dim, 1)
+
+    def forward(self, state):
+        features = self.feature_layer(state)
+
+        # Actor输出动作概率
+        action_probs = F.softmax(self.actor(features), dim=-1)
+
+        # Critic输出状态价值
+        state_value = self.critic(features)
+
+        return action_probs, state_value
+
+
+class ActorCriticAgent:
+    """基础的Actor-Critic算法"""
+
+    def __init__(self,
+                 env,
+                 actor_lr=1e-3,
+                 critic_lr=1e-3,
+                 gamma=0.99,
+                 hidden_dim=256,
+                 render=False):
+
+        self.env = env
+        self.gamma = gamma
+        self.render = render
+
+        self.state_dim = env.observation_space.shape[0]
+        self.action_dim = env.action_space.n
+
+        # 使用共享网络或分离网络
+        self.use_shared_network = True
+
+        if self.use_shared_network:
+            # 共享网络版本
+            self.actor_critic = ActorCritic(self.state_dim, self.action_dim, hidden_dim)
+            self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=actor_lr)
+        else:
+            # 分离网络版本
+            self.actor = nn.Sequential(
+                nn.Linear(self.state_dim, hidden_dim),
+                nn.ReLU(),
+                nn.Linear(hidden_dim, hidden_dim),
+                nn.ReLU(),
+                nn.Linear(hidden_dim, self.action_dim),
+                nn.Softmax(dim=-1)
+            )
+            self.critic = nn.Sequential(
+                nn.Linear(self.state_dim, hidden_dim),
+                nn.ReLU(),
+                nn.Linear(hidden_dim, hidden_dim),
+                nn.ReLU(),
+                nn.Linear(hidden_dim, 1)
+            )
+            self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=actor_lr)
+            self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=critic_lr)
+
+        # 训练统计
+        self.training_stats = {
+            'episode_rewards': [],
+            'critic_loss': [],
+            'actor_loss': [],
+            'advantages': []
+        }
+
+    def select_action(self, state):
+        """选择动作并返回动作、log概率和状态价值"""
+        state = torch.FloatTensor(state).unsqueeze(0)
+
+        if self.use_shared_network:
+            action_probs, state_value = self.actor_critic(state)
+        else:
+            action_probs = self.actor(state)
+            state_value = self.critic(state)
+
+        m = Categorical(action_probs)
+        action = m.sample()
+        log_prob = m.log_prob(action)
+
+        return action.item(), log_prob, state_value
+
+    def update(self, log_prob, state_value, reward, next_state_value, done):
+        """单步更新Actor和Critic"""
+        if self.use_shared_network:
+            return self._update_shared(log_prob, state_value, reward, next_state_value, done)
+        else:
+            return self._update_separate(log_prob, state_value, reward, next_state_value, done)
+
+    def _update_shared(self, log_prob, state_value, reward, next_state_value, done):
+        """共享网络更新"""
+        # 计算TD目标
+        td_target = reward + (1 - done) * self.gamma * next_state_value
+        td_target = td_target.detach()
+
+        # 计算TD误差（优势）
+        td_error = td_target - state_value
+
+        # Critic损失
+        critic_loss = td_error.pow(2).mean()
+
+        # Actor损失
+        actor_loss = -(log_prob * td_error.detach()).mean()
+        # td_error，实际的未来奖励 - 预期的未来奖励
+        # 如果是正数，则加大当前动作的概率，如果是负数则减小动作的概率。
+
+        # 总损失
+        total_loss = actor_loss + 0.5 * critic_loss
+
+        # 反向传播
+        self.optimizer.zero_grad()
+        total_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.actor_critic.parameters(), 0.5)
+        self.optimizer.step()
+
+        return actor_loss.item(), critic_loss.item(), td_error.mean().item()
+
+    def _update_separate(self, log_prob, state_value, reward, next_state_value, done):
+        """分离网络更新"""
+        # 计算TD目标
+        td_target = reward + (1 - done) * self.gamma * next_state_value
+        td_target = td_target.detach()
+
+        # 计算TD误差
+        td_error = td_target - state_value
+
+        # 更新Critic
+        critic_loss = td_error.pow(2).mean()
+        self.critic_optimizer.zero_grad()
+        critic_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 0.5)
+        self.critic_optimizer.step()
+
+        # 更新Actor
+        actor_loss = -(log_prob * td_error.detach()).mean()
+        self.actor_optimizer.zero_grad()
+        actor_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 0.5)
+        self.actor_optimizer.step()
+
+        return actor_loss.item(), critic_loss.item(), td_error.mean().item()
+
+    def train(self, num_episodes=1000, max_steps=5000):
+        """训练智能体"""
+        episode_rewards = []
+        episode_lengths = []
+
+        for episode in range(num_episodes):
+            state, _ = self.env.reset()
+            episode_reward = 0
+            episode_losses = {'actor': [], 'critic': [], 'advantage': []}
+
+            for step in range(max_steps):
+                if self.render:
+                    self.env.render()
+
+                # 选择动作
+                action, log_prob, state_value = self.select_action(state)
+
+                # 执行动作
+                next_state, reward, terminated, truncated, _ = self.env.step(action)
+                done = terminated or truncated
+
+                # 获取下一状态的价值
+                with torch.no_grad():
+                    next_state_tensor = torch.FloatTensor(next_state).unsqueeze(0)
+                    if self.use_shared_network:
+                        _, next_state_value = self.actor_critic(next_state_tensor)
+                    else:
+                        next_state_value = self.critic(next_state_tensor)
+
+                # 更新网络
+                actor_loss, critic_loss, advantage = self.update(
+                    log_prob, state_value, reward, next_state_value, done
+                )
+
+                episode_reward += reward
+                episode_losses['actor'].append(actor_loss)
+                episode_losses['critic'].append(critic_loss)
+                episode_losses['advantage'].append(advantage)
+
+                state = next_state
+
+                if done:
+                    break
+
+            # 记录统计信息
+            episode_rewards.append(episode_reward)
+            episode_lengths.append(step + 1)
+            self.training_stats['episode_rewards'].append(episode_reward)
+            self.training_stats['actor_loss'].append(np.mean(episode_losses['actor']))
+            self.training_stats['critic_loss'].append(np.mean(episode_losses['critic']))
+            self.training_stats['advantages'].append(np.mean(episode_losses['advantage']))
+
+            # 打印进度
+            if (episode + 1) % 20 == 0:
+                avg_reward = np.mean(episode_rewards[-20:])
+                avg_actor_loss = np.mean(self.training_stats['actor_loss'][-20:])
+                avg_critic_loss = np.mean(self.training_stats['critic_loss'][-20:])
+                print(f"回合 {episode + 1:4d} | "
+                      f"奖励: {episode_reward:5.1f} | "
+                      f"平均奖励: {avg_reward:5.1f} | "
+                      f"A-Loss: {avg_actor_loss:.4f} | "
+                      f"C-Loss: {avg_critic_loss:.4f}")
+
+        return episode_rewards
+
+
+# ============== 2. A2C (Advantage Actor-Critic) ==============
+class A2C(nn.Module):
+    """A2C网络 - 包含熵正则化"""
+
+    def __init__(self, state_dim, action_dim, hidden_dim=256):
+        super(A2C, self).__init__()
+
+        self.actor = nn.Sequential(
+            nn.Linear(state_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, action_dim)
+        )
+
+        self.critic = nn.Sequential(
+            nn.Linear(state_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 1)
+        )
+
+    def forward(self, state):
+        action_logits = self.actor(state)
+        state_value = self.critic(state)
+        return action_logits, state_value
+
+
+class A2CAgent:
+    """A2C算法 - 使用优势函数和熵正则化"""
+
+    def __init__(self,
+                 env,
+                 learning_rate=3e-4,
+                 gamma=0.99,
+                 gae_lambda=0.95,
+                 entropy_coef=0.01,
+                 value_coef=0.5,
+                 max_grad_norm=0.5,
+                 hidden_dim=256):
+
+        self.env = env
+        self.gamma = gamma
+        self.gae_lambda = gae_lambda
+        self.entropy_coef = entropy_coef
+        self.value_coef = value_coef
+        self.max_grad_norm = max_grad_norm
+
+        self.state_dim = env.observation_space.shape[0]
+        self.action_dim = env.action_space.n
+
+        self.network = A2C(self.state_dim, self.action_dim, hidden_dim)
+        self.optimizer = optim.Adam(self.network.parameters(), lr=learning_rate)
+
+        # 存储经验
+        self.states = []
+        self.actions = []
+        self.log_probs = []
+        self.rewards = []
+        self.values = []
+        self.dones = []
+
+        self.training_stats = {
+            'episode_rewards': [],
+            'policy_loss': [],
+            'value_loss': [],
+            'entropy': []
+        }
+
+    def select_action(self, state):
+        """选择动作并存储经验"""
+        state = torch.FloatTensor(state).unsqueeze(0)
+        action_logits, state_value = self.network(state)
+
+        action_probs = F.softmax(action_logits, dim=-1)
+        m = Categorical(action_probs)
+        action = m.sample()
+        log_prob = m.log_prob(action)
+
+        # 存储经验
+        self.states.append(state)
+        self.actions.append(action)
+        self.log_probs.append(log_prob)
+        self.values.append(state_value)
+
+        return action.item()
+
+    def compute_gae(self, rewards, values, dones):
+        """计算广义优势估计(GAE)"""
+        advantages = []
+        gae = 0
+
+        values = values + [0]  # 添加最后一个虚拟价值
+
+        for t in reversed(range(len(rewards))):
+            delta = rewards[t] + self.gamma * values[t + 1] * (1 - dones[t]) - values[t]
+            gae = delta + self.gamma * self.gae_lambda * (1 - dones[t]) * gae
+            advantages.insert(0, gae)
+
+        returns = [adv + val for adv, val in zip(advantages, values[:-1])]
+
+        return advantages, returns
+
+    def update(self):
+        """更新网络参数"""
+        if len(self.rewards) == 0:
+            return
+
+        # 转换为tensor
+        states = torch.cat(self.states)
+        actions = torch.cat(self.actions)
+        old_log_probs = torch.cat(self.log_probs).detach()
+        rewards = self.rewards
+        values = [v.squeeze() for v in self.values]
+        dones = self.dones
+
+        # 计算GAE和returns
+        advantages, returns = self.compute_gae(rewards, values, dones)
+        advantages = torch.FloatTensor(advantages)
+        returns = torch.FloatTensor(returns)
+
+        # 标准化优势
+        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+
+        # 重新计算log probs和values
+        action_logits, state_values = self.network(states)
+        action_probs = F.softmax(action_logits, dim=-1)
+        m = Categorical(action_probs)
+        log_probs = m.log_prob(actions)
+        entropy = m.entropy().mean()
+
+        # 计算损失
+        state_values = state_values.squeeze()
+        value_loss = F.mse_loss(state_values, returns)
+
+        policy_loss = -(log_probs * advantages.detach()).mean()
+        entropy_loss = -self.entropy_coef * entropy
+
+        total_loss = policy_loss + self.value_coef * value_loss + entropy_loss
+
+        # 反向传播
+        self.optimizer.zero_grad()
+        total_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.network.parameters(), self.max_grad_norm)
+        self.optimizer.step()
+
+        # 记录统计
+        policy_loss_val = policy_loss.item()
+        value_loss_val = value_loss.item()
+        entropy_val = entropy.item()
+
+        # 清空经验
+        self.states = []
+        self.actions = []
+        self.log_probs = []
+        self.rewards = []
+        self.values = []
+        self.dones = []
+
+        return policy_loss_val, value_loss_val, entropy_val
+
+    def train(self, num_episodes=1000, max_steps=500, update_frequency=10):
+        """训练A2C智能体"""
+        episode_rewards = []
+        episode_lengths = []
+
+        for episode in range(num_episodes):
+            state, _ = self.env.reset()
+            episode_reward = 0
+
+            for step in range(max_steps):
+                # 选择动作
+                action = self.select_action(state)
+
+                # 执行动作
+                next_state, reward, terminated, truncated, _ = self.env.step(action)
+                done = terminated or truncated
+
+                # 存储经验
+                self.rewards.append(reward)
+                self.dones.append(done)
+
+                episode_reward += reward
+                state = next_state
+
+                if done:
+                    break
+
+            episode_rewards.append(episode_reward)
+            episode_lengths.append(step + 1)
+            self.training_stats['episode_rewards'].append(episode_reward)
+
+            # 更新网络
+            if (episode + 1) % update_frequency == 0:
+                policy_loss, value_loss, entropy = self.update()
+                self.training_stats['policy_loss'].append(policy_loss)
+                self.training_stats['value_loss'].append(value_loss)
+                self.training_stats['entropy'].append(entropy)
+
+            # 打印进度
+            if (episode + 1) % 20 == 0:
+                avg_reward = np.mean(episode_rewards[-20:])
+                print(f"A2C - 回合 {episode + 1:4d} | "
+                      f"奖励: {episode_reward:5.1f} | "
+                      f"平均奖励: {avg_reward:5.1f}")
+
+        return episode_rewards
+
+
+# ============== 3. 可视化对比 ==============
+def compare_algorithms():
+    """对比基础Actor-Critic和A2C"""
+    print("\n" + "=" * 70)
+    print("Actor-Critic算法对比实验")
+    print("=" * 70)
+
+    # 创建环境
+    env = gym.make('CartPole-v1')
+
+    # 1. 基础Actor-Critic
+    print("\n1. 训练基础Actor-Critic...")
+    ac_agent = ActorCriticAgent(env, actor_lr=1e-3, critic_lr=1e-3)
+    ac_rewards = ac_agent.train(num_episodes=300)
+
+    # 2. A2C
+    print("\n2. 训练A2C (带GAE和熵正则化)...")
+    a2c_agent = A2CAgent(env, learning_rate=3e-4)
+    a2c_rewards = a2c_agent.train(num_episodes=300)
+
+    # 可视化对比
+    fig, axes = plt.subplots(2, 3, figsize=(16, 10))
+
+    # 1. 奖励曲线对比
+    ax1 = axes[0, 0]
+    ax1.plot(ac_rewards, alpha=0.6, label='Actor-Critic', color='blue')
+    ax1.plot(a2c_rewards, alpha=0.6, label='A2C', color='red')
+
+    # 平滑曲线
+    window = 20
+    ac_smooth = np.convolve(ac_rewards, np.ones(window) / window, mode='valid')
+    a2c_smooth = np.convolve(a2c_rewards, np.ones(window) / window, mode='valid')
+
+    ax1.plot(range(window - 1, len(ac_smooth) + window - 1), ac_smooth,
+             'b-', linewidth=2, label='AC (平滑)')
+    ax1.plot(range(window - 1, len(a2c_smooth) + window - 1), a2c_smooth,
+             'r-', linewidth=2, label='A2C (平滑)')
+
+    ax1.set_xlabel('回合')
+    ax1.set_ylabel('总奖励')
+    ax1.set_title('训练奖励对比')
+    ax1.legend()
+    ax1.grid(True, alpha=0.3)
+
+    # 2. Actor损失对比
+    ax2 = axes[0, 1]
+    if hasattr(ac_agent.training_stats, 'actor_loss'):
+        ax2.plot(ac_agent.training_stats['actor_loss'],
+                 label='Actor-Critic', color='blue', alpha=0.7)
+    ax2.plot(a2c_agent.training_stats['policy_loss'],
+             label='A2C', color='red', alpha=0.7)
+    ax2.set_xlabel('更新步')
+    ax2.set_ylabel('策略损失')
+    ax2.set_title('Actor损失对比')
+    ax2.legend()
+    ax2.grid(True, alpha=0.3)
+
+    # 3. Critic损失对比
+    ax3 = axes[0, 2]
+    if hasattr(ac_agent.training_stats, 'critic_loss'):
+        ax3.plot(ac_agent.training_stats['critic_loss'],
+                 label='Actor-Critic', color='blue', alpha=0.7)
+    ax3.plot(a2c_agent.training_stats['value_loss'],
+             label='A2C', color='red', alpha=0.7)
+    ax3.set_xlabel('更新步')
+    ax3.set_ylabel('价值损失')
+    ax3.set_title('Critic损失对比')
+    ax3.legend()
+    ax3.grid(True, alpha=0.3)
+
+    # 4. 熵变化 (A2C特有)
+    ax4 = axes[1, 0]
+    if a2c_agent.training_stats['entropy']:
+        ax4.plot(a2c_agent.training_stats['entropy'],
+                 color='green', linewidth=2)
+    ax4.set_xlabel('更新步')
+    ax4.set_ylabel('策略熵')
+    ax4.set_title('A2C - 策略熵变化')
+    ax4.grid(True, alpha=0.3)
+
+    # 5. 收敛速度箱线图
+    ax5 = axes[1, 1]
+
+    # 计算收敛回合数
+    def get_convergence_episode(rewards, threshold=450):
+        for i, r in enumerate(rewards):
+            if r >= threshold and np.mean(rewards[max(0, i - 10):i + 1]) >= threshold:
+                return i + 1
+        return len(rewards)
+
+    # 多次实验
+    n_trials = 10
+    ac_convergence = []
+    a2c_convergence = []
+
+    for trial in range(n_trials):
+        env_trial = gym.make('CartPole-v1')
+
+        ac_trial = ActorCriticAgent(env_trial, actor_lr=1e-3, critic_lr=1e-3)
+        ac_rewards_trial = ac_trial.train(num_episodes=200)
+        ac_convergence.append(get_convergence_episode(ac_rewards_trial))
+
+        a2c_trial = A2CAgent(env_trial, learning_rate=3e-4)
+        a2c_rewards_trial = a2c_trial.train(num_episodes=200)
+        a2c_convergence.append(get_convergence_episode(a2c_rewards_trial))
+
+    bp = ax5.boxplot([ac_convergence, a2c_convergence],
+                     labels=['Actor-Critic', 'A2C'],
+                     patch_artist=True)
+    bp['boxes'][0].set_facecolor('lightblue')
+    bp['boxes'][1].set_facecolor('lightcoral')
+    ax5.set_ylabel('收敛所需回合数')
+    ax5.set_title('收敛速度对比 (越低越好)')
+    ax5.grid(True, alpha=0.3)
+
+    # 6. 算法对比表格
+    ax6 = axes[1, 2]
+    ax6.axis('off')
+
+    # 创建对比表格
+    col_labels = ['算法', '更新方式', '优势估计', '熵正则', '收敛速度', '稳定性']
+    data = [
+        ['Actor-Critic', '单步TD', 'TD误差', '无', '较慢', '中等'],
+        ['A2C', '多步回报', 'GAE', '有', '快', '稳定']
+    ]
+
+    table = ax6.table(cellText=data,
+                      colLabels=col_labels,
+                      cellLoc='center',
+                      loc='center',
+                      bbox=[0, 0, 1, 1])
+    table.auto_set_font_size(False)
+    table.set_fontsize(10)
+    table.scale(1, 1.5)
+
+    plt.suptitle('Actor-Critic vs A2C: 性能对比分析', fontsize=14, y=1.02)
+    plt.tight_layout()
+    plt.savefig('ac_vs_a2c_comparison.png', dpi=150, bbox_inches='tight')
+    plt.show()
+
+    # 打印总结
+    print("\n" + "=" * 70)
+    print("实验总结:")
+    print("=" * 70)
+    print(f"\n基础Actor-Critic:")
+    print(f"  - 平均收敛回合: {np.mean(ac_convergence):.1f} ± {np.std(ac_convergence):.1f}")
+    print(f"  - 最终平均奖励: {np.mean(ac_rewards[-50:]):.1f}")
+
+    print(f"\nA2C (带GAE和熵正则):")
+    print(f"  - 平均收敛回合: {np.mean(a2c_convergence):.1f} ± {np.std(a2c_convergence):.1f}")
+    print(f"  - 最终平均奖励: {np.mean(a2c_rewards[-50:]):.1f}")
+
+    improvement = (np.mean(ac_convergence) - np.mean(a2c_convergence)) / np.mean(ac_convergence) * 100
+    print(f"\nA2C收敛速度提升: {improvement:.1f}%")
+
+    return ac_agent, a2c_agent
+
+
+# ============== 4. 主函数 ==============
+def main():
+    """主函数"""
+    print("=" * 70)
+    print("Actor-Critic 方法完整实现")
+    print("=" * 70)
+
+    # 参数设置
+    import argparse
+    parser = argparse.ArgumentParser(description='Actor-Critic算法实现')
+    parser.add_argument('--algo', type=str, default='a2c',
+                        choices=['ac', 'a2c', 'compare'],
+                        help='选择算法: ac (基础Actor-Critic), a2c, compare')
+    parser.add_argument('--episodes', type=int, default=5000,
+                        help='训练回合数')
+    parser.add_argument('--render', action='store_true',
+                        help='渲染环境')
+    args = parser.parse_args()
+
+    env = gym.make('CartPole-v1')
+
+    if args.algo == 'ac':
+        print("\n训练基础Actor-Critic...")
+        agent = ActorCriticAgent(env, render=args.render)
+        rewards = agent.train(num_episodes=args.episodes)
+
+        plt.figure(figsize=(10, 6))
+        plt.plot(rewards, alpha=0.6, label='Episode Reward')
+
+        # 平滑曲线
+        window = 20
+        if len(rewards) >= window:
+            smoothed = np.convolve(rewards, np.ones(window) / window, mode='valid')
+            plt.plot(range(window - 1, len(smoothed) + window - 1),
+                     smoothed, 'r-', linewidth=2, label=f'Moving Avg (window={window})')
+
+        plt.xlabel('Episode')
+        plt.ylabel('Total Reward')
+        plt.title('Actor-Critic Training on CartPole')
+        plt.legend()
+        plt.grid(True, alpha=0.3)
+        plt.savefig('actor_critic_training.png', dpi=150)
+        plt.show()
+
+    elif args.algo == 'a2c':
+        print("\n训练A2C...")
+        agent = A2CAgent(env, learning_rate=3e-4)
+        rewards = agent.train(num_episodes=args.episodes)
+
+        fig, axes = plt.subplots(2, 2, figsize=(12, 8))
+
+        # 奖励曲线
+        axes[0, 0].plot(rewards, alpha=0.6, color='blue')
+        if len(rewards) >= 20:
+            smoothed = np.convolve(rewards, np.ones(20) / 20, mode='valid')
+            axes[0, 0].plot(range(19, len(smoothed) + 19), smoothed, 'r-', linewidth=2)
+        axes[0, 0].set_xlabel('Episode')
+        axes[0, 0].set_ylabel('Total Reward')
+        axes[0, 0].set_title('A2C Training Rewards')
+        axes[0, 0].grid(True, alpha=0.3)
+
+        # 策略损失
+        axes[0, 1].plot(agent.training_stats['policy_loss'], color='purple')
+        axes[0, 1].set_xlabel('Update Step')
+        axes[0, 1].set_ylabel('Policy Loss')
+        axes[0, 1].set_title('Policy Loss')
+        axes[0, 1].grid(True, alpha=0.3)
+
+        # 价值损失
+        axes[1, 0].plot(agent.training_stats['value_loss'], color='orange')
+        axes[1, 0].set_xlabel('Update Step')
+        axes[1, 0].set_ylabel('Value Loss')
+        axes[1, 0].set_title('Value Loss')
+        axes[1, 0].grid(True, alpha=0.3)
+
+        # 策略熵
+        axes[1, 1].plot(agent.training_stats['entropy'], color='green')
+        axes[1, 1].set_xlabel('Update Step')
+        axes[1, 1].set_ylabel('Policy Entropy')
+        axes[1, 1].set_title('Policy Entropy')
+        axes[1, 1].grid(True, alpha=0.3)
+
+        plt.tight_layout()
+        plt.savefig('a2c_training.png', dpi=150)
+        plt.show()
+
+    else:  # compare
+        compare_algorithms()
+
+    env.close()
+    print("\n训练完成!")
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rl/cart_pole/step_2_ppo_clip.py

@@ -0,0 +1,739 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+🌱 **第一代：REINFORCE (1992)**
+├── 核心创新：蒙特卡洛策略梯度
+├── 公式：∇θ ∝ Σ ∇θ log π * Gₜ
+├── 痛点：▸ 必须等回合结束 ▸ 方差极大 ▸ 样本效率低
+└── 贡献：开创了策略梯度范式
+
+🌿 **第二代：REINFORCE with Baseline (2000s)**
+├── 核心创新：引入基线降低方差
+├── 公式：∇θ ∝ Σ ∇θ log π * (Gₜ - b(s))
+├── 痛点：▸ 仍需完整回合 ▸ 基线需要单独学习
+└── 贡献：方差降低，训练更稳定
+
+🍃 **第三代：Actor-Critic (2000s)**
+├── 核心创新：单步更新，不再等待回合
+├── 公式：∇θ ∝ ∇θ log π * (r + γV(s') - V(s))
+├── 痛点：▸ 单步TD偏差大 ▸ 价值估计不准
+└── 贡献：实现了真正的在线学习
+
+🌳 **第四代：A2C/A3C (2016)**
+├── 核心创新：优势函数 + 多步回报
+├── 公式：∇θ ∝ ∇θ log π * Â(s,a)
+├── GAE：Â = Σ (γλ)ᵏ δ_{t+k}  ★滑动平均★
+├── 痛点：▸ 更新步长敏感 ▸ 容易破坏策略
+└── 贡献：GAE成为标准配置
+
+🌲 **第五代：PPO (2017)**
+├── 核心创新：**所有前人智慧的集大成**
+├── 1️⃣ 继承AC/A2C：单步更新 + GAE
+├── 2️⃣ 继承重要性采样：可以复用数据
+├── 3️⃣ ✨ **独创：Clipped Surrogate Objective**
+│      L = min(r(θ)Â, clip(r(θ), 1-ε, 1+ε)Â)
+├── 4️⃣ ✨ **独创：自适应KL惩罚**
+└── 贡献：**稳定、高效、易用，成为事实标准**
+
+由于函数的优化步长受到优势的直接影响。当优步长过大时很容易直接跨过，导致优化失败。
+虽然A2C 已经对历史优势进行了移动平均，但问题仍然存在。
+尤其是当训练的早期价值函数还没有获得较好的训练，这种问题尤其容易出现。
+
+clip 当新模型相比于旧模型动作的概率变化幅度 radio 已经较大时，则不再进行优化。
+clip 操作会对 radio 值进行截断，会切断梯度反向传播的通道。
+
+同时对动作的概率做熵最大化，目的是避免策略函数过早确定化，保持探索能力。
+
+PPO-clip 使用radio * advantages
+其中：ratio = probs / old_probs
+而不是 A2C 中的 log_probs * advantages
+
+
+"""
+import gymnasium as gym
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.distributions import Categorical
+import numpy as np
+import matplotlib.pyplot as plt
+from collections import deque
+import time
+
+
+# ============== 1. PPO-Clip Network ==============
+class PPONetwork(nn.Module):
+    """PPO Network: Shared feature extractor + Actor head + Critic head"""
+
+    def __init__(self, state_dim, action_dim, hidden_dim=64):
+        super(PPONetwork, self).__init__()
+
+        # Shared feature extractor
+        self.feature_layer = nn.Sequential(
+            nn.Linear(state_dim, hidden_dim),
+            nn.Tanh(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.Tanh()
+        )
+
+        # Actor head: action probability distribution
+        self.actor = nn.Sequential(
+            nn.Linear(hidden_dim, action_dim),
+            nn.Softmax(dim=-1)
+        )
+
+        # Critic head: state value
+        self.critic = nn.Linear(hidden_dim, 1)
+
+    def forward(self, state):
+        features = self.feature_layer(state)
+        action_probs = self.actor(features)
+        state_value = self.critic(features)
+        return action_probs, state_value
+
+
+# ============== 2. PPO-Clip Agent ==============
+class PPOClipAgent:
+    """PPO-Clip Agent with Gymnasium API"""
+
+    def __init__(self,
+                 env,
+                 learning_rate=3e-4,
+                 gamma=0.99,
+                 gae_lambda=0.95,
+                 clip_epsilon=0.2,
+                 entropy_coef=0.01,
+                 value_coef=0.5,
+                 max_grad_norm=0.5,
+                 update_epochs=4,
+                 mini_batch_size=64,
+                 horizon=2048,
+                 hidden_dim=64):
+
+        self.env = env
+        self.gamma = gamma
+        self.gae_lambda = gae_lambda
+        self.clip_epsilon = clip_epsilon
+        self.entropy_coef = entropy_coef
+        self.value_coef = value_coef
+        self.max_grad_norm = max_grad_norm
+        self.update_epochs = update_epochs
+        self.mini_batch_size = mini_batch_size
+        self.horizon = horizon
+
+        self.state_dim = env.observation_space.shape[0]
+        self.action_dim = env.action_space.n
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.network = PPONetwork(self.state_dim, self.action_dim, hidden_dim).to(self.device)
+        self.optimizer = optim.Adam(self.network.parameters(), lr=learning_rate)
+
+        self.reset_buffer()
+
+        self.training_stats = {
+            'episode_rewards': [],
+            'episode_lengths': [],
+            'policy_loss': [],
+            'value_loss': [],
+            'entropy': [],
+            'clip_fraction': [],
+            'explained_variance': []
+        }
+
+        # Statistics for table logging
+        self.recent_rewards = deque(maxlen=20)
+        self.recent_lengths = deque(maxlen=20)
+
+    def reset_buffer(self):
+        """Reset experience buffer"""
+        self.buffer = {
+            'states': [],
+            'actions': [],
+            'rewards': [],
+            'next_states': [],
+            'dones': [],
+            'log_probs': [],
+            'values': []
+        }
+
+    def select_action(self, state, eval_mode=False):
+        """Select action using current policy"""
+        state = torch.FloatTensor(state).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            action_probs, state_value = self.network(state)
+
+        m = Categorical(action_probs)
+        action = m.sample()
+        log_prob = m.log_prob(action)
+
+        if eval_mode:
+            action = torch.argmax(action_probs)
+            log_prob = m.log_prob(action)
+
+        return action.item(), log_prob.cpu().item(), state_value.cpu().item()
+
+    def store_transition(self, state, action, reward, next_state, done, log_prob, value):
+        """Store one step of experience"""
+        self.buffer['states'].append(state)
+        self.buffer['actions'].append(action)
+        self.buffer['rewards'].append(reward)
+        self.buffer['next_states'].append(next_state)
+        self.buffer['dones'].append(done)
+        self.buffer['log_probs'].append(log_prob)
+        self.buffer['values'].append(value)
+
+    def compute_gae(self, rewards, values, next_values, dones):
+        """Compute Generalized Advantage Estimation"""
+        advantages = []
+        gae = 0
+
+        for t in reversed(range(len(rewards))):
+            delta = rewards[t] + self.gamma * next_values[t] * (1 - dones[t]) - values[t]
+            gae = delta + self.gamma * self.gae_lambda * (1 - dones[t]) * gae
+            advantages.insert(0, gae)
+
+        return advantages
+
+    def update(self):
+        """PPO-Clip update"""
+        if len(self.buffer['rewards']) < self.mini_batch_size:
+            return None
+
+        # Convert to tensors
+        states = torch.FloatTensor(np.array(self.buffer['states'])).to(self.device)
+        actions = torch.LongTensor(self.buffer['actions']).to(self.device)
+        old_log_probs = torch.FloatTensor(self.buffer['log_probs']).to(self.device)
+
+        values = torch.FloatTensor(self.buffer['values']).to(self.device)
+        next_states = torch.FloatTensor(np.array(self.buffer['next_states'])).to(self.device)
+        dones = torch.FloatTensor(self.buffer['dones']).to(self.device)
+        rewards = self.buffer['rewards']
+
+        # Compute next state values
+        with torch.no_grad():
+            _, next_values = self.network(next_states)
+            next_values = next_values.squeeze().cpu().numpy()
+
+        # Compute advantages and returns
+        advantages = self.compute_gae(
+            rewards,
+            values.cpu().numpy(),
+            next_values,
+            dones.cpu().numpy()
+        )
+        advantages = torch.FloatTensor(advantages).to(self.device)
+        returns = advantages + values
+
+        # Normalize advantages
+        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+
+        total_policy_loss = 0
+        total_value_loss = 0
+        total_entropy = 0
+        total_clip_fraction = 0
+
+        dataset_size = len(states)
+
+        # Multiple epochs of PPO update
+        for epoch in range(self.update_epochs):
+            indices = np.random.permutation(dataset_size)
+
+            for start in range(0, dataset_size, self.mini_batch_size):
+                end = start + self.mini_batch_size
+                batch_indices = indices[start:end]
+
+                batch_states = states[batch_indices]
+                batch_actions = actions[batch_indices]
+                batch_old_log_probs = old_log_probs[batch_indices]
+                batch_advantages = advantages[batch_indices]
+                batch_returns = returns[batch_indices]
+
+                # Forward pass
+                action_probs, state_values = self.network(batch_states)
+                state_values = state_values.squeeze()
+
+                # Compute log probs and entropy
+                m = Categorical(action_probs)
+                log_probs = m.log_prob(batch_actions)
+                entropy = m.entropy().mean()
+
+                # Importance sampling ratio
+                ratio = torch.exp(log_probs - batch_old_log_probs)
+
+                # PPO-Clip objective
+                surr1 = ratio * batch_advantages
+                surr2 = torch.clamp(ratio, 1 - self.clip_epsilon, 1 + self.clip_epsilon) * batch_advantages
+                policy_loss = -torch.min(surr1, surr2).mean()
+
+                # Value loss
+                value_loss = nn.MSELoss()(state_values, batch_returns)
+
+                # Entropy loss
+                entropy_loss = -self.entropy_coef * entropy
+
+                # Total loss
+                total_loss = policy_loss + self.value_coef * value_loss + entropy_loss
+
+                # Backward pass
+                self.optimizer.zero_grad()
+                total_loss.backward()
+                torch.nn.utils.clip_grad_norm_(self.network.parameters(), self.max_grad_norm)
+                self.optimizer.step()
+
+                total_policy_loss += policy_loss.item()
+                total_value_loss += value_loss.item()
+                total_entropy += entropy.item()
+
+                # Compute clip fraction
+                with torch.no_grad():
+                    clip_mask = (ratio < 1 - self.clip_epsilon) | (ratio > 1 + self.clip_epsilon)
+                    clip_fraction = clip_mask.float().mean().item()
+                    total_clip_fraction += clip_fraction
+
+        n_updates = self.update_epochs * (dataset_size // self.mini_batch_size + 1)
+
+        # Compute explained variance
+        with torch.no_grad():
+            _, pred_values = self.network(states)
+            pred_values = pred_values.squeeze()
+            explained_variance = 1 - torch.var(returns - pred_values) / (torch.var(returns) + 1e-8)
+            explained_variance = explained_variance.cpu().item()
+
+        stats = {
+            'policy_loss': total_policy_loss / n_updates,
+            'value_loss': total_value_loss / n_updates,
+            'entropy': total_entropy / n_updates,
+            'clip_fraction': total_clip_fraction / n_updates,
+            'explained_variance': explained_variance
+        }
+
+        self.reset_buffer()
+
+        return stats
+
+    def print_header(self):
+        """Print table header"""
+        print("\n" + "=" * 130)
+        print(
+            f"{'Episode':>8} | {'Avg Reward':>12} | {'Avg Length':>12} | {'Policy Loss':>12} | {'Value Loss':>12} | {'Entropy':>10} | {'Clip%':>10} | {'Expl Var':>12} | {'Time':>10}")
+        print("-" * 130)
+
+    def print_row(self, episode, total_episodes, stats=None, elapsed_time=None):
+        """Print table row"""
+
+        # Update statistics cache
+        if len(self.training_stats['episode_rewards']) > 0:
+            self.recent_rewards.append(self.training_stats['episode_rewards'][-1])
+        if len(self.training_stats['episode_lengths']) > 0:
+            self.recent_lengths.append(self.training_stats['episode_lengths'][-1])
+
+        # Calculate averages
+        avg_reward = np.mean(self.recent_rewards) if self.recent_rewards else 0
+        avg_length = np.mean(self.recent_lengths) if self.recent_lengths else 0
+
+        # Format output
+        if stats:
+            print(f"{episode:>8}/{total_episodes} | "
+                  f"{avg_reward:>12.2f} | "
+                  f"{avg_length:>12.1f} | "
+                  f"{stats['policy_loss']:>12.4f} | "
+                  f"{stats['value_loss']:>12.4f} | "
+                  f"{stats['entropy']:>10.4f} | "
+                  f"{stats['clip_fraction']:>10.3f} | "
+                  f"{stats['explained_variance']:>12.3f} | "
+                  f"{elapsed_time:>10.1f}s")
+        else:
+            print(f"{episode:>8}/{total_episodes} | "
+                  f"{avg_reward:>12.2f} | "
+                  f"{avg_length:>12.1f} | "
+                  f"{'-':>12} | "
+                  f"{'-':>12} | "
+                  f"{'-':>10} | "
+                  f"{'-':>10} | "
+                  f"{'-':>12} | "
+                  f"{elapsed_time:>10.1f}s")
+
+    def print_summary(self, total_time, num_episodes):
+        """Print training summary"""
+        print("=" * 130)
+        print(f"\n🎯 Training completed! Episodes: {num_episodes}, Total time: {total_time:.1f}s")
+
+        if len(self.training_stats['episode_rewards']) >= 20:
+            final_avg_reward = np.mean(self.training_stats['episode_rewards'][-20:])
+            final_avg_length = np.mean(self.training_stats['episode_lengths'][-20:])
+            print(f"📊 Last 20 episodes - Avg Reward: {final_avg_reward:.2f}, Avg Length: {final_avg_length:.1f}")
+
+        if self.training_stats['policy_loss']:
+            print(f"📉 Final Policy Loss: {self.training_stats['policy_loss'][-1]:.4f}")
+            print(f"📉 Final Value Loss: {self.training_stats['value_loss'][-1]:.4f}")
+            print(f"🎲 Final Entropy: {self.training_stats['entropy'][-1]:.4f}")
+            print(f"✂️  Final Clip Fraction: {self.training_stats['clip_fraction'][-1]:.3f}")
+            print(f"📈 Final Explained Variance: {self.training_stats['explained_variance'][-1]:.3f}")
+
+        print("=" * 130)
+
+    def train(self, num_episodes=1000, max_steps_per_episode=500, log_interval=20):
+        """Train PPO-Clip agent"""
+
+        print("\n" + "🚀" * 65)
+        print("PPO-Clip Training Started (Gymnasium API)")
+        print("🚀" * 65)
+
+        print(f"\n📋 Hyperparameters:")
+        print(f"   Learning Rate: {self.optimizer.param_groups[0]['lr']:.6f}")
+        print(f"   Gamma: {self.gamma:.2f}, GAE Lambda: {self.gae_lambda:.2f}")
+        print(f"   Clip Epsilon: {self.clip_epsilon:.2f}, Update Epochs: {self.update_epochs}")
+        print(f"   Mini-batch: {self.mini_batch_size}, Horizon: {self.horizon}")
+        print(f"   Entropy Coef: {self.entropy_coef:.3f}, Value Coef: {self.value_coef:.1f}")
+        print(f"   Device: {self.device}")
+
+        self.print_header()
+
+        total_steps = 0
+        episode = 0
+        start_time = time.time()
+
+        while episode < num_episodes:
+            # Gymnasium API: reset() returns state, info
+            state, _ = self.env.reset()
+            episode_reward = 0
+            episode_step = 0
+
+            while episode_step < max_steps_per_episode:
+                action, log_prob, value = self.select_action(state)
+
+                # Gymnasium API: step() returns next_state, reward, terminated, truncated, info
+                next_state, reward, terminated, truncated, _ = self.env.step(action)
+                done = terminated or truncated
+
+                self.store_transition(state, action, reward, next_state, done, log_prob, value)
+
+                episode_reward += reward
+                episode_step += 1
+                total_steps += 1
+                state = next_state
+
+                # Update when buffer is full or episode ends
+                if len(self.buffer['rewards']) >= self.horizon or done:
+                    stats = self.update()
+                    if stats:
+                        self.training_stats['policy_loss'].append(stats['policy_loss'])
+                        self.training_stats['value_loss'].append(stats['value_loss'])
+                        self.training_stats['entropy'].append(stats['entropy'])
+                        self.training_stats['clip_fraction'].append(stats['clip_fraction'])
+                        self.training_stats['explained_variance'].append(stats['explained_variance'])
+
+                if done:
+                    break
+
+            # Record episode statistics
+            self.training_stats['episode_rewards'].append(episode_reward)
+            self.training_stats['episode_lengths'].append(episode_step)
+
+            episode += 1
+
+            # Log periodically
+            if episode % log_interval == 0:
+                current_time = time.time()
+                elapsed = current_time - start_time
+
+                recent_stats = None
+                if self.training_stats['policy_loss']:
+                    recent_stats = {
+                        'policy_loss': self.training_stats['policy_loss'][-1],
+                        'value_loss': self.training_stats['value_loss'][-1],
+                        'entropy': self.training_stats['entropy'][-1],
+                        'clip_fraction': self.training_stats['clip_fraction'][-1],
+                        'explained_variance': self.training_stats['explained_variance'][-1]
+                    }
+
+                self.print_row(episode, num_episodes, recent_stats, elapsed)
+
+        total_time = time.time() - start_time
+        self.print_summary(total_time, num_episodes)
+
+        return self.training_stats['episode_rewards'], self.training_stats['episode_lengths']
+
+    def save(self, path):
+        """Save model"""
+        torch.save({
+            'network_state_dict': self.network.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'training_stats': self.training_stats
+        }, path)
+        print(f"\n💾 Model saved to {path}")
+
+    def load(self, path):
+        """Load model"""
+        checkpoint = torch.load(path)
+        self.network.load_state_dict(checkpoint['network_state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        self.training_stats = checkpoint['training_stats']
+        print(f"\n📂 Model loaded from {path}")
+
+
+# ============== 3. Evaluation Function ==============
+def evaluate_agent(agent, env, num_episodes=10, render=False):
+    """Evaluate trained agent"""
+    print("\n" + "🎯" * 35)
+    print("Evaluation Started")
+    print("🎯" * 35)
+
+    print(f"\n{'Episode':^12} | {'Reward':^12} | {'Length':^12} | {'Avg Reward':^14}")
+    print("-" * 60)
+
+    episode_rewards = []
+    episode_lengths = []
+
+    for episode in range(num_episodes):
+        state, _ = env.reset()
+        episode_reward = 0
+        episode_step = 0
+
+        while True:
+            if render:
+                env.render()
+                time.sleep(0.02)
+
+            action, _, _ = agent.select_action(state, eval_mode=True)
+            next_state, reward, terminated, truncated, _ = env.step(action)
+            done = terminated or truncated
+
+            episode_reward += reward
+            episode_step += 1
+            state = next_state
+
+            if done:
+                break
+
+        episode_rewards.append(episode_reward)
+        episode_lengths.append(episode_step)
+
+        avg_so_far = np.mean(episode_rewards)
+        print(f"{episode + 1:^12} | {episode_reward:^12.1f} | {episode_step:^12} | {avg_so_far:^14.2f}")
+
+    print("-" * 60)
+    print(f"\n📊 Evaluation Results ({num_episodes} episodes):")
+    print(f"   Avg Reward: {np.mean(episode_rewards):.2f} ± {np.std(episode_rewards):.2f}")
+    print(f"   Avg Length: {np.mean(episode_lengths):.2f} ± {np.std(episode_lengths):.2f}")
+    print(f"   Max Reward: {np.max(episode_rewards):.2f}")
+    print(f"   Min Reward: {np.min(episode_rewards):.2f}")
+    print(f"   Success Rate (>=475): {np.mean(np.array(episode_rewards) >= 475) * 100:.1f}%")
+    print("=" * 60)
+
+    return episode_rewards, episode_lengths
+
+
+# ============== 4. Visualization Function (English Only) ==============
+def plot_training_results(agent, save_path='ppo_training_results.png'):
+    """Plot training results - English only, no Chinese font issues"""
+
+    stats = agent.training_stats
+
+    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
+    fig.suptitle('PPO-Clip Training Results (CartPole-v1)', fontsize=16, y=1.02)
+
+    window = 20
+
+    # 1. Episode Rewards
+    ax1 = axes[0, 0]
+    rewards = stats['episode_rewards']
+    ax1.plot(rewards, alpha=0.3, color='blue', label='Raw Reward')
+
+    if len(rewards) >= window:
+        smoothed = np.convolve(rewards, np.ones(window) / window, mode='valid')
+        ax1.plot(range(window - 1, len(smoothed) + window - 1), smoothed,
+                 'r-', linewidth=2, label=f'{window}-Episode MA')
+
+    ax1.axhline(y=500, color='green', linestyle='--', alpha=0.7, label='Target (500)')
+    ax1.set_xlabel('Episode')
+    ax1.set_ylabel('Total Reward')
+    ax1.set_title('Training Rewards')
+    ax1.legend()
+    ax1.grid(True, alpha=0.3)
+
+    # 2. Episode Lengths
+    ax2 = axes[0, 1]
+    lengths = stats['episode_lengths']
+    ax2.plot(lengths, alpha=0.3, color='orange', label='Episode Length')
+
+    if len(lengths) >= window:
+        smoothed_length = np.convolve(lengths, np.ones(window) / window, mode='valid')
+        ax2.plot(range(window - 1, len(smoothed_length) + window - 1), smoothed_length,
+                 'r-', linewidth=2)
+
+    ax2.set_xlabel('Episode')
+    ax2.set_ylabel('Episode Length')
+    ax2.set_title('Episode Lengths')
+    ax2.grid(True, alpha=0.3)
+
+    # 3. Policy Loss
+    ax3 = axes[0, 2]
+    if stats['policy_loss']:
+        ax3.plot(stats['policy_loss'], color='purple', linewidth=1.5)
+        ax3.set_xlabel('Update Step')
+        ax3.set_ylabel('Policy Loss')
+        ax3.set_title('Policy Loss')
+        ax3.grid(True, alpha=0.3)
+
+    # 4. Value Loss
+    ax4 = axes[1, 0]
+    if stats['value_loss']:
+        ax4.plot(stats['value_loss'], color='brown', linewidth=1.5)
+        ax4.set_xlabel('Update Step')
+        ax4.set_ylabel('Value Loss')
+        ax4.set_title('Value Loss')
+        ax4.grid(True, alpha=0.3)
+
+    # 5. Policy Entropy
+    ax5 = axes[1, 1]
+    if stats['entropy']:
+        ax5.plot(stats['entropy'], color='green', linewidth=1.5)
+        ax5.set_xlabel('Update Step')
+        ax5.set_ylabel('Policy Entropy')
+        ax5.set_title('Policy Entropy (Exploration)')
+        ax5.grid(True, alpha=0.3)
+
+    # 6. Clip Fraction & Explained Variance
+    ax6 = axes[1, 2]
+    if stats['clip_fraction']:
+        ax6.plot(stats['clip_fraction'], color='red', linewidth=1.5, label='Clip Fraction')
+        ax6.set_xlabel('Update Step')
+        ax6.set_ylabel('Clip Fraction', color='red')
+        ax6.tick_params(axis='y', labelcolor='red')
+        ax6.grid(True, alpha=0.3)
+
+        ax6_twin = ax6.twinx()
+        if stats['explained_variance']:
+            ax6_twin.plot(stats['explained_variance'], color='blue', linewidth=1.5,
+                          label='Explained Variance')
+            ax6_twin.set_ylabel('Explained Variance', color='blue')
+            ax6_twin.tick_params(axis='y', labelcolor='blue')
+
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=150, bbox_inches='tight')
+    plt.show()
+    print(f"\n📸 Training results saved to {save_path}")
+
+
+# ============== 5. Main Function ==============
+def main():
+    """Main function"""
+
+    # Create environment (Gymnasium)
+    env = gym.make('CartPole-v1')
+
+    # PPO-Clip hyperparameters
+    config = {
+        'learning_rate': 3e-4,
+        'gamma': 0.99,
+        'gae_lambda': 0.95,
+        'clip_epsilon': 0.2,
+        'entropy_coef': 0.01,
+        'value_coef': 0.5,
+        'max_grad_norm': 0.5,
+        'update_epochs': 4,
+        'mini_batch_size': 64,
+        'horizon': 2048,
+        'hidden_dim': 64
+    }
+
+    print("\n" + "=" * 90)
+    print("PPO-Clip for CartPole-v1 (Gymnasium)")
+    print("=" * 90)
+    print("\n📋 Hyperparameters:")
+    for key, value in config.items():
+        print(f"   {key:20}: {value}")
+
+    # Create PPO agent
+    agent = PPOClipAgent(env, **config)
+
+    try:
+        # Train
+        rewards, lengths = agent.train(num_episodes=5000, log_interval=20)
+
+        # Save model
+        agent.save('ppo_cartpole_gymnasium.pth')
+
+        # Plot results
+        plot_training_results(agent)
+
+        # Evaluate
+        print("\n")
+        eval_env = gym.make('CartPole-v1')
+        evaluate_agent(agent, eval_env, num_episodes=20, render=False)
+
+        # Demo (with rendering)
+        print("\n")
+        demo_env = gym.make('CartPole-v1', render_mode='human')
+        evaluate_agent(agent, demo_env, num_episodes=3, render=True)
+
+    except KeyboardInterrupt:
+        print("\n\n⚠️ Training interrupted, saving model...")
+        agent.save('ppo_cartpole_gymnasium_interrupted.pth')
+
+    finally:
+        env.close()
+        if 'eval_env' in locals():
+            eval_env.close()
+        if 'demo_env' in locals():
+            demo_env.close()
+
+
+# ============== 6. Hyperparameter Sweep ==============
+def hyperparameter_sweep():
+    """Hyperparameter tuning experiment"""
+
+    print("\n" + "🔬" * 45)
+    print("PPO-Clip Hyperparameter Sweep")
+    print("🔬" * 45)
+
+    clip_values = [0.1, 0.2, 0.3]
+    results = {}
+
+    for clip_eps in clip_values:
+        print(f"\n📊 Testing clip_epsilon = {clip_eps}")
+        print("-" * 60)
+
+        env = gym.make('CartPole-v1')
+        agent = PPOClipAgent(
+            env,
+            learning_rate=3e-4,
+            clip_epsilon=clip_eps,
+            update_epochs=4,
+            horizon=2048,
+            mini_batch_size=64
+        )
+
+        rewards, _ = agent.train(num_episodes=2000, log_interval=20)
+        results[f'ε={clip_eps}'] = rewards
+        env.close()
+
+    # Plot comparison
+    plt.figure(figsize=(12, 6))
+
+    for name, rewards in results.items():
+        window = 20
+        smoothed = np.convolve(rewards, np.ones(window) / window, mode='valid')
+        plt.plot(range(window - 1, len(smoothed) + window - 1), smoothed,
+                 linewidth=2, label=name)
+
+    plt.xlabel('Episode')
+    plt.ylabel(f'Avg Reward ({window}-Episode MA)')
+    plt.title('PPO-Clip: Different Clip Epsilon Comparison')
+    plt.legend()
+    plt.grid(True, alpha=0.3)
+    plt.savefig('ppo_clip_comparison.png', dpi=150)
+    plt.show()
+
+    return results
+
+
+if __name__ == "__main__":
+    main()
+    # hyperparameter_sweep()  # Uncomment to run hyperparameter sweep
+
+
+

examples/tutorials/rl/cart_pole/step_2_ppo_penalty.py

@@ -0,0 +1,767 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+KL散度损失就是为了直接约束新旧策略之间的变化程度。
+
+使用 KL散度的好处：
+penalty 更像是，不管优势多大，它总能将其进行可控的相对缩放。
+
+
+clip
+当新模型相比旧模型的变化幅度已较大时，clip以阻断优化，会切断梯度传导。
+penalty
+直接将新模型与旧模型的动作概率约束在一个 target_kl 附近，限制每一次迭代的优化幅度。
+KL散度不会切断梯度传导，总是可以进行有效的优化。
+
+"""
+import gymnasium as gym
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.distributions import Categorical
+import numpy as np
+import matplotlib.pyplot as plt
+from collections import deque
+import time
+
+
+# ============== 1. PPO-Penalty Network ==============
+class PPONetwork(nn.Module):
+    """PPO Network: Shared feature extractor + Actor head + Critic head"""
+
+    def __init__(self, state_dim, action_dim, hidden_dim=64):
+        super(PPONetwork, self).__init__()
+
+        # Shared feature extractor
+        self.feature_layer = nn.Sequential(
+            nn.Linear(state_dim, hidden_dim),
+            nn.Tanh(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.Tanh()
+        )
+
+        # Actor head: action probability distribution
+        self.actor = nn.Sequential(
+            nn.Linear(hidden_dim, action_dim),
+            nn.Softmax(dim=-1)
+        )
+
+        # Critic head: state value
+        self.critic = nn.Linear(hidden_dim, 1)
+
+    def forward(self, state):
+        features = self.feature_layer(state)
+        action_probs = self.actor(features)
+        state_value = self.critic(features)
+        return action_probs, state_value
+
+
+# ============== 2. PPO-Penalty Agent ==============
+class PPOPenaltyAgent:
+    """PPO-Penalty (Adaptive KL Penalty) with Gymnasium API"""
+
+    def __init__(self,
+                 env,
+                 learning_rate=3e-4,
+                 gamma=0.99,
+                 gae_lambda=0.95,
+                 kl_target=0.01,  # Target KL divergence
+                 kl_coef_init=1.0,  # Initial KL penalty coefficient
+                 kl_coef_adapt=1.5,  # KL coefficient adaptation rate
+                 entropy_coef=0.01,
+                 value_coef=0.5,
+                 max_grad_norm=0.5,
+                 update_epochs=10,  # PPO-Penalty typically uses more epochs
+                 mini_batch_size=64,
+                 horizon=2048,
+                 hidden_dim=64):
+
+        self.env = env
+        self.gamma = gamma
+        self.gae_lambda = gae_lambda
+        self.kl_target = kl_target
+        self.kl_coef = kl_coef_init
+        self.kl_coef_adapt = kl_coef_adapt
+        self.entropy_coef = entropy_coef
+        self.value_coef = value_coef
+        self.max_grad_norm = max_grad_norm
+        self.update_epochs = update_epochs
+        self.mini_batch_size = mini_batch_size
+        self.horizon = horizon
+
+        self.state_dim = env.observation_space.shape[0]
+        self.action_dim = env.action_space.n
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Policy network
+        self.policy = PPONetwork(self.state_dim, self.action_dim, hidden_dim).to(self.device)
+        self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate)
+
+        # Old policy for KL calculation
+        self.old_policy = PPONetwork(self.state_dim, self.action_dim, hidden_dim).to(self.device)
+        self.update_old_policy()
+
+        self.reset_buffer()
+
+        self.training_stats = {
+            'episode_rewards': [],
+            'episode_lengths': [],
+            'policy_loss': [],
+            'value_loss': [],
+            'entropy': [],
+            'kl_divergence': [],
+            'kl_coef': [],
+            'explained_variance': []
+        }
+
+        # Statistics for table logging
+        self.recent_rewards = deque(maxlen=20)
+        self.recent_lengths = deque(maxlen=20)
+
+    def update_old_policy(self):
+        """Copy current policy to old policy"""
+        self.old_policy.load_state_dict(self.policy.state_dict())
+
+    def reset_buffer(self):
+        """Reset experience buffer"""
+        self.buffer = {
+            'states': [],
+            'actions': [],
+            'rewards': [],
+            'next_states': [],
+            'dones': [],
+            'log_probs': [],
+            'values': []
+        }
+
+    def select_action(self, state, eval_mode=False):
+        """Select action using current policy"""
+        state = torch.FloatTensor(state).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            action_probs, state_value = self.policy(state)
+
+        m = Categorical(action_probs)
+        action = m.sample()
+        log_prob = m.log_prob(action)
+
+        if eval_mode:
+            action = torch.argmax(action_probs)
+            log_prob = m.log_prob(action)
+
+        return action.item(), log_prob.cpu().item(), state_value.cpu().item()
+
+    def store_transition(self, state, action, reward, next_state, done, log_prob, value):
+        """Store one step of experience"""
+        self.buffer['states'].append(state)
+        self.buffer['actions'].append(action)
+        self.buffer['rewards'].append(reward)
+        self.buffer['next_states'].append(next_state)
+        self.buffer['dones'].append(done)
+        self.buffer['log_probs'].append(log_prob)
+        self.buffer['values'].append(value)
+
+    def compute_gae(self, rewards, values, next_values, dones):
+        """Compute Generalized Advantage Estimation"""
+        advantages = []
+        gae = 0
+
+        for t in reversed(range(len(rewards))):
+            delta = rewards[t] + self.gamma * next_values[t] * (1 - dones[t]) - values[t]
+            gae = delta + self.gamma * self.gae_lambda * (1 - dones[t]) * gae
+            advantages.insert(0, gae)
+
+        return advantages
+
+    def compute_kl_divergence(self, states, actions):
+        """Compute KL divergence between old and new policy"""
+        with torch.no_grad():
+            # Get old policy distributions
+            old_probs, _ = self.old_policy(states)
+            old_m = Categorical(old_probs)
+
+            # Get new policy distributions
+            new_probs, _ = self.policy(states)
+            new_m = Categorical(new_probs)
+
+            # Compute KL divergence
+            kl = torch.distributions.kl.kl_divergence(old_m, new_m).mean()
+
+        return kl.item()
+
+    def update(self):
+        """PPO-Penalty update with adaptive KL penalty"""
+        if len(self.buffer['rewards']) < self.mini_batch_size:
+            return None
+
+        # Convert to tensors
+        states = torch.FloatTensor(np.array(self.buffer['states'])).to(self.device)
+        actions = torch.LongTensor(self.buffer['actions']).to(self.device)
+        old_log_probs = torch.FloatTensor(self.buffer['log_probs']).to(self.device)
+
+        values = torch.FloatTensor(self.buffer['values']).to(self.device)
+        next_states = torch.FloatTensor(np.array(self.buffer['next_states'])).to(self.device)
+        dones = torch.FloatTensor(self.buffer['dones']).to(self.device)
+        rewards = self.buffer['rewards']
+
+        # Compute next state values
+        with torch.no_grad():
+            _, next_values = self.policy(next_states)
+            next_values = next_values.squeeze().cpu().numpy()
+
+        # Compute advantages and returns
+        advantages = self.compute_gae(
+            rewards,
+            values.cpu().numpy(),
+            next_values,
+            dones.cpu().numpy()
+        )
+        advantages = torch.FloatTensor(advantages).to(self.device)
+        returns = advantages + values
+
+        # Normalize advantages
+        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+
+        total_policy_loss = 0
+        total_value_loss = 0
+        total_entropy = 0
+        total_kl = 0
+
+        dataset_size = len(states)
+
+        # Multiple epochs of PPO update
+        for epoch in range(self.update_epochs):
+            indices = np.random.permutation(dataset_size)
+
+            for start in range(0, dataset_size, self.mini_batch_size):
+                end = start + self.mini_batch_size
+                batch_indices = indices[start:end]
+
+                batch_states = states[batch_indices]
+                batch_actions = actions[batch_indices]
+                batch_old_log_probs = old_log_probs[batch_indices]
+                batch_advantages = advantages[batch_indices]
+                batch_returns = returns[batch_indices]
+
+                # Forward pass
+                action_probs, state_values = self.policy(batch_states)
+                state_values = state_values.squeeze()
+
+                # Compute log probs and entropy
+                m = Categorical(action_probs)
+                log_probs = m.log_prob(batch_actions)
+                entropy = m.entropy().mean()
+
+                # Importance sampling ratio
+                ratio = torch.exp(log_probs - batch_old_log_probs)
+
+                # PPO-Penalty objective (with KL penalty, no clipping!)
+                policy_loss = -(ratio * batch_advantages).mean()
+
+                # Compute KL divergence for this batch
+                with torch.no_grad():
+                    old_probs, _ = self.old_policy(batch_states)
+                    old_m = Categorical(old_probs)
+                    kl_batch = torch.distributions.kl.kl_divergence(old_m, m).mean()
+                    total_kl += kl_batch.item()
+
+                # Add KL penalty to policy loss
+                policy_loss_penalized = policy_loss + self.kl_coef * kl_batch
+
+                # Value loss
+                value_loss = nn.MSELoss()(state_values, batch_returns)
+
+                # Entropy loss (encourage exploration)
+                entropy_loss = -self.entropy_coef * entropy
+
+                # Total loss
+                total_loss = policy_loss_penalized + self.value_coef * value_loss + entropy_loss
+
+                # Backward pass
+                self.optimizer.zero_grad()
+                total_loss.backward()
+                torch.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
+                self.optimizer.step()
+
+                total_policy_loss += policy_loss.item()
+                total_value_loss += value_loss.item()
+                total_entropy += entropy.item()
+
+        # Compute average KL divergence
+        n_updates = self.update_epochs * (dataset_size // self.mini_batch_size + 1)
+        avg_kl = total_kl / n_updates
+
+        # Adapt KL coefficient (核心：自适应调整KL惩罚系数)
+        if avg_kl < self.kl_target / 1.5:
+            # KL too small -> reduce penalty
+            self.kl_coef /= self.kl_coef_adapt
+        elif avg_kl > self.kl_target * 1.5:
+            # KL too large -> increase penalty
+            self.kl_coef *= self.kl_coef_adapt
+
+        # Keep KL coefficient in reasonable range
+        self.kl_coef = np.clip(self.kl_coef, 1e-10, 10.0)
+
+        # Compute explained variance
+        with torch.no_grad():
+            _, pred_values = self.policy(states)
+            pred_values = pred_values.squeeze()
+            explained_variance = 1 - torch.var(returns - pred_values) / (torch.var(returns) + 1e-8)
+            explained_variance = explained_variance.cpu().item()
+
+        stats = {
+            'policy_loss': total_policy_loss / n_updates,
+            'value_loss': total_value_loss / n_updates,
+            'entropy': total_entropy / n_updates,
+            'kl_divergence': avg_kl,
+            'kl_coef': self.kl_coef,
+            'explained_variance': explained_variance
+        }
+
+        # Update old policy
+        self.update_old_policy()
+        self.reset_buffer()
+
+        return stats
+
+    def print_header(self):
+        """Print table header"""
+        print("\n" + "=" * 150)
+        print(
+            f"{'Episode':>8} | {'Avg Reward':>12} | {'Avg Length':>12} | {'Policy Loss':>12} | {'Value Loss':>12} | {'Entropy':>10} | {'KL Div':>10} | {'KL Coef':>10} | {'Expl Var':>12} | {'Time':>10}")
+        print("-" * 150)
+
+    def print_row(self, episode, total_episodes, stats=None, elapsed_time=None):
+        """Print table row"""
+
+        # Update statistics cache
+        if len(self.training_stats['episode_rewards']) > 0:
+            self.recent_rewards.append(self.training_stats['episode_rewards'][-1])
+        if len(self.training_stats['episode_lengths']) > 0:
+            self.recent_lengths.append(self.training_stats['episode_lengths'][-1])
+
+        # Calculate averages
+        avg_reward = np.mean(self.recent_rewards) if self.recent_rewards else 0
+        avg_length = np.mean(self.recent_lengths) if self.recent_lengths else 0
+
+        # Format output
+        if stats:
+            print(f"{episode:>8}/{total_episodes} | "
+                  f"{avg_reward:>12.2f} | "
+                  f"{avg_length:>12.1f} | "
+                  f"{stats['policy_loss']:>12.4f} | "
+                  f"{stats['value_loss']:>12.4f} | "
+                  f"{stats['entropy']:>10.4f} | "
+                  f"{stats['kl_divergence']:>10.6f} | "
+                  f"{stats['kl_coef']:>10.6f} | "
+                  f"{stats['explained_variance']:>12.3f} | "
+                  f"{elapsed_time:>10.1f}s")
+        else:
+            print(f"{episode:>8}/{total_episodes} | "
+                  f"{avg_reward:>12.2f} | "
+                  f"{avg_length:>12.1f} | "
+                  f"{'-':>12} | "
+                  f"{'-':>12} | "
+                  f"{'-':>10} | "
+                  f"{'-':>10} | "
+                  f"{'-':>10} | "
+                  f"{'-':>12} | "
+                  f"{elapsed_time:>10.1f}s")
+
+    def print_summary(self, total_time, num_episodes):
+        """Print training summary"""
+        print("=" * 150)
+        print(f"\n🎯 Training completed! Episodes: {num_episodes}, Total time: {total_time:.1f}s")
+
+        if len(self.training_stats['episode_rewards']) >= 20:
+            final_avg_reward = np.mean(self.training_stats['episode_rewards'][-20:])
+            final_avg_length = np.mean(self.training_stats['episode_lengths'][-20:])
+            print(f"📊 Last 20 episodes - Avg Reward: {final_avg_reward:.2f}, Avg Length: {final_avg_length:.1f}")
+
+        if self.training_stats['policy_loss']:
+            print(f"📉 Final Policy Loss: {self.training_stats['policy_loss'][-1]:.4f}")
+            print(f"📉 Final Value Loss: {self.training_stats['value_loss'][-1]:.4f}")
+            print(f"🎲 Final Entropy: {self.training_stats['entropy'][-1]:.4f}")
+            print(f"📏 Final KL Divergence: {self.training_stats['kl_divergence'][-1]:.6f}")
+            print(f"⚖️  Final KL Coefficient: {self.training_stats['kl_coef'][-1]:.6f}")
+            print(f"📈 Final Explained Variance: {self.training_stats['explained_variance'][-1]:.3f}")
+
+        print("=" * 150)
+
+    def train(self, num_episodes=1000, max_steps_per_episode=500, log_interval=20):
+        """Train PPO-Penalty agent"""
+
+        print("\n" + "🚀" * 75)
+        print("PPO-Penalty (Adaptive KL) Training Started - Gymnasium API")
+        print("🚀" * 75)
+
+        print(f"\n📋 Hyperparameters:")
+        print(f"   Learning Rate: {self.optimizer.param_groups[0]['lr']:.6f}")
+        print(f"   Gamma: {self.gamma:.2f}, GAE Lambda: {self.gae_lambda:.2f}")
+        print(f"   KL Target: {self.kl_target:.4f}, KL Coef Init: {self.kl_coef:.3f}")
+        print(f"   KL Adapt Rate: {self.kl_coef_adapt:.2f}")
+        print(f"   Update Epochs: {self.update_epochs}")
+        print(f"   Mini-batch: {self.mini_batch_size}, Horizon: {self.horizon}")
+        print(f"   Entropy Coef: {self.entropy_coef:.3f}, Value Coef: {self.value_coef:.1f}")
+        print(f"   Device: {self.device}")
+
+        self.print_header()
+
+        total_steps = 0
+        episode = 0
+        start_time = time.time()
+
+        while episode < num_episodes:
+            # Gymnasium API: reset() returns state, info
+            state, _ = self.env.reset()
+            episode_reward = 0
+            episode_step = 0
+
+            while episode_step < max_steps_per_episode:
+                action, log_prob, value = self.select_action(state)
+
+                # Gymnasium API: step() returns next_state, reward, terminated, truncated, info
+                next_state, reward, terminated, truncated, _ = self.env.step(action)
+                done = terminated or truncated
+
+                self.store_transition(state, action, reward, next_state, done, log_prob, value)
+
+                episode_reward += reward
+                episode_step += 1
+                total_steps += 1
+                state = next_state
+
+                # Update when buffer is full or episode ends
+                if len(self.buffer['rewards']) >= self.horizon or done:
+                    stats = self.update()
+                    if stats:
+                        self.training_stats['policy_loss'].append(stats['policy_loss'])
+                        self.training_stats['value_loss'].append(stats['value_loss'])
+                        self.training_stats['entropy'].append(stats['entropy'])
+                        self.training_stats['kl_divergence'].append(stats['kl_divergence'])
+                        self.training_stats['kl_coef'].append(stats['kl_coef'])
+                        self.training_stats['explained_variance'].append(stats['explained_variance'])
+
+                if done:
+                    break
+
+            # Record episode statistics
+            self.training_stats['episode_rewards'].append(episode_reward)
+            self.training_stats['episode_lengths'].append(episode_step)
+
+            episode += 1
+
+            # Log periodically
+            if episode % log_interval == 0:
+                current_time = time.time()
+                elapsed = current_time - start_time
+
+                recent_stats = None
+                if self.training_stats['policy_loss']:
+                    recent_stats = {
+                        'policy_loss': self.training_stats['policy_loss'][-1],
+                        'value_loss': self.training_stats['value_loss'][-1],
+                        'entropy': self.training_stats['entropy'][-1],
+                        'kl_divergence': self.training_stats['kl_divergence'][-1],
+                        'kl_coef': self.training_stats['kl_coef'][-1],
+                        'explained_variance': self.training_stats['explained_variance'][-1]
+                    }
+
+                self.print_row(episode, num_episodes, recent_stats, elapsed)
+
+        total_time = time.time() - start_time
+        self.print_summary(total_time, num_episodes)
+
+        return self.training_stats['episode_rewards'], self.training_stats['episode_lengths']
+
+    def save(self, path):
+        """Save model"""
+        torch.save({
+            'policy_state_dict': self.policy.state_dict(),
+            'old_policy_state_dict': self.old_policy.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'kl_coef': self.kl_coef,
+            'training_stats': self.training_stats
+        }, path)
+        print(f"\n💾 Model saved to {path}")
+
+    def load(self, path):
+        """Load model"""
+        checkpoint = torch.load(path)
+        self.policy.load_state_dict(checkpoint['policy_state_dict'])
+        self.old_policy.load_state_dict(checkpoint['old_policy_state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        self.kl_coef = checkpoint['kl_coef']
+        self.training_stats = checkpoint['training_stats']
+        print(f"\n📂 Model loaded from {path}")
+
+
+# ============== 3. Evaluation Function ==============
+def evaluate_agent(agent, env, num_episodes=10, render=False):
+    """Evaluate trained agent"""
+    print("\n" + "🎯" * 35)
+    print("Evaluation Started")
+    print("🎯" * 35)
+
+    print(f"\n{'Episode':^12} | {'Reward':^12} | {'Length':^12} | {'Avg Reward':^14}")
+    print("-" * 60)
+
+    episode_rewards = []
+    episode_lengths = []
+
+    for episode in range(num_episodes):
+        state, _ = env.reset()
+        episode_reward = 0
+        episode_step = 0
+
+        while True:
+            if render:
+                env.render()
+                time.sleep(0.02)
+
+            action, _, _ = agent.select_action(state, eval_mode=True)
+            next_state, reward, terminated, truncated, _ = env.step(action)
+            done = terminated or truncated
+
+            episode_reward += reward
+            episode_step += 1
+            state = next_state
+
+            if done:
+                break
+
+        episode_rewards.append(episode_reward)
+        episode_lengths.append(episode_step)
+
+        avg_so_far = np.mean(episode_rewards)
+        print(f"{episode + 1:^12} | {episode_reward:^12.1f} | {episode_step:^12} | {avg_so_far:^14.2f}")
+
+    print("-" * 60)
+    print(f"\n📊 Evaluation Results ({num_episodes} episodes):")
+    print(f"   Avg Reward: {np.mean(episode_rewards):.2f} ± {np.std(episode_rewards):.2f}")
+    print(f"   Avg Length: {np.mean(episode_lengths):.2f} ± {np.std(episode_lengths):.2f}")
+    print(f"   Max Reward: {np.max(episode_rewards):.2f}")
+    print(f"   Min Reward: {np.min(episode_rewards):.2f}")
+    print(f"   Success Rate (>=475): {np.mean(np.array(episode_rewards) >= 475) * 100:.1f}%")
+    print("=" * 60)
+
+    return episode_rewards, episode_lengths
+
+
+# ============== 4. Visualization Function ==============
+def plot_training_results(agent, save_path='ppo_penalty_training_results.png'):
+    """Plot training results"""
+
+    stats = agent.training_stats
+
+    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
+    fig.suptitle('PPO-Penalty (Adaptive KL) Training Results - CartPole-v1', fontsize=16, y=1.02)
+
+    window = 20
+
+    # 1. Episode Rewards
+    ax1 = axes[0, 0]
+    rewards = stats['episode_rewards']
+    ax1.plot(rewards, alpha=0.3, color='blue', label='Raw Reward')
+
+    if len(rewards) >= window:
+        smoothed = np.convolve(rewards, np.ones(window) / window, mode='valid')
+        ax1.plot(range(window - 1, len(smoothed) + window - 1), smoothed,
+                 'r-', linewidth=2, label=f'{window}-Episode MA')
+
+    ax1.axhline(y=500, color='green', linestyle='--', alpha=0.7, label='Target (500)')
+    ax1.set_xlabel('Episode')
+    ax1.set_ylabel('Total Reward')
+    ax1.set_title('Training Rewards')
+    ax1.legend()
+    ax1.grid(True, alpha=0.3)
+
+    # 2. Episode Lengths
+    ax2 = axes[0, 1]
+    lengths = stats['episode_lengths']
+    ax2.plot(lengths, alpha=0.3, color='orange', label='Episode Length')
+
+    if len(lengths) >= window:
+        smoothed_length = np.convolve(lengths, np.ones(window) / window, mode='valid')
+        ax2.plot(range(window - 1, len(smoothed_length) + window - 1), smoothed_length,
+                 'r-', linewidth=2)
+
+    ax2.set_xlabel('Episode')
+    ax2.set_ylabel('Episode Length')
+    ax2.set_title('Episode Lengths')
+    ax2.grid(True, alpha=0.3)
+
+    # 3. Policy Loss & KL Divergence
+    ax3 = axes[0, 2]
+    if stats['policy_loss']:
+        ax3_twin = ax3.twinx()
+        ax3.plot(stats['policy_loss'], color='purple', linewidth=1.5, label='Policy Loss')
+        ax3.set_xlabel('Update Step')
+        ax3.set_ylabel('Policy Loss', color='purple')
+        ax3.tick_params(axis='y', labelcolor='purple')
+
+        if stats['kl_divergence']:
+            ax3_twin.plot(stats['kl_divergence'], color='orange', linewidth=1.5, label='KL Div')
+            ax3_twin.set_ylabel('KL Divergence', color='orange')
+            ax3_twin.tick_params(axis='y', labelcolor='orange')
+    ax3.grid(True, alpha=0.3)
+
+    # 4. Value Loss
+    ax4 = axes[1, 0]
+    if stats['value_loss']:
+        ax4.plot(stats['value_loss'], color='brown', linewidth=1.5)
+        ax4.set_xlabel('Update Step')
+        ax4.set_ylabel('Value Loss')
+        ax4.set_title('Value Loss')
+        ax4.grid(True, alpha=0.3)
+
+    # 5. Policy Entropy
+    ax5 = axes[1, 1]
+    if stats['entropy']:
+        ax5.plot(stats['entropy'], color='green', linewidth=1.5)
+        ax5.set_xlabel('Update Step')
+        ax5.set_ylabel('Policy Entropy')
+        ax5.set_title('Policy Entropy (Exploration)')
+        ax5.grid(True, alpha=0.3)
+
+    # 6. KL Coefficient & Explained Variance
+    ax6 = axes[1, 2]
+    if stats['kl_coef']:
+        ax6_twin = ax6.twinx()
+        ax6.plot(stats['kl_coef'], color='red', linewidth=1.5, label='KL Coef')
+        ax6.set_xlabel('Update Step')
+        ax6.set_ylabel('KL Coefficient', color='red')
+        ax6.tick_params(axis='y', labelcolor='red')
+
+        if stats['explained_variance']:
+            ax6_twin.plot(stats['explained_variance'], color='blue', linewidth=1.5, label='Expl Var')
+            ax6_twin.set_ylabel('Explained Variance', color='blue')
+            ax6_twin.tick_params(axis='y', labelcolor='blue')
+    ax6.grid(True, alpha=0.3)
+
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=150, bbox_inches='tight')
+    plt.show()
+    print(f"\n📸 Training results saved to {save_path}")
+
+
+# ============== 5. Main Function ==============
+def main():
+    """Main function"""
+
+    # Create environment (Gymnasium)
+    env = gym.make('CartPole-v1')
+
+    # PPO-Penalty hyperparameters
+    config = {
+        'learning_rate': 3e-4,
+        'gamma': 0.99,
+        'gae_lambda': 0.95,
+        'kl_target': 0.01,  # Target KL divergence per update
+        'kl_coef_init': 1.0,  # Initial KL penalty coefficient
+        'kl_coef_adapt': 1.5,  # Adaptation rate
+        'entropy_coef': 0.01,
+        'value_coef': 0.5,
+        'max_grad_norm': 0.5,
+        'update_epochs': 10,  # More epochs for PPO-Penalty
+        'mini_batch_size': 64,
+        'horizon': 2048,
+        'hidden_dim': 64
+    }
+
+    print("\n" + "=" * 100)
+    print("PPO-Penalty (Adaptive KL) for CartPole-v1 (Gymnasium)")
+    print("=" * 100)
+    print("\n📋 Hyperparameters:")
+    for key, value in config.items():
+        print(f"   {key:20}: {value}")
+
+    # Create PPO-Penalty agent
+    agent = PPOPenaltyAgent(env, **config)
+
+    try:
+        # Train
+        rewards, lengths = agent.train(num_episodes=500, log_interval=20)
+
+        # Save model
+        agent.save('ppo_penalty_cartpole.pth')
+
+        # Plot results
+        plot_training_results(agent)
+
+        # Evaluate
+        print("\n")
+        eval_env = gym.make('CartPole-v1')
+        evaluate_agent(agent, eval_env, num_episodes=20, render=False)
+
+        # Demo (with rendering)
+        print("\n")
+        demo_env = gym.make('CartPole-v1', render_mode='human')
+        evaluate_agent(agent, demo_env, num_episodes=3, render=True)
+
+    except KeyboardInterrupt:
+        print("\n\n⚠️ Training interrupted, saving model...")
+        agent.save('ppo_penalty_cartpole_interrupted.pth')
+
+    finally:
+        env.close()
+        if 'eval_env' in locals():
+            eval_env.close()
+        if 'demo_env' in locals():
+            demo_env.close()
+
+
+# ============== 6. Compare PPO-Clip vs PPO-Penalty ==============
+def compare_ppo_variants():
+    """Compare PPO-Clip and PPO-Penalty"""
+
+    print("\n" + "🔬" * 50)
+    print("PPO-Clip vs PPO-Penalty Comparison")
+    print("🔬" * 50)
+
+    # This would require implementing both agents and running experiments
+    # For brevity, here's the conceptual comparison:
+
+    comparison = """
+    📊 **PPO-Clip vs PPO-Penalty Comparison**
+
+    ============================================================
+    Feature           | PPO-Clip              | PPO-Penalty
+    ============================================================
+    Constraint Type   | Hard clipping         | Soft KL penalty
+    ------------------------------------------------------------
+    Update Limit      | r ∈ [1-ε, 1+ε]        | KL(π||π_old) < target
+    ------------------------------------------------------------
+    Adaptation        | Fixed ε               | Adaptive KL coef
+    ------------------------------------------------------------
+    Implementation    | Simple                | More complex
+    ------------------------------------------------------------
+    Compute Cost      | Low                   | Higher (KL calc)
+    ------------------------------------------------------------
+    Stability         | Very stable           | Very stable
+    ------------------------------------------------------------
+    Sample Efficiency | Good                  | Good
+    ------------------------------------------------------------
+    Hyperparameter    | ε=0.2 (robust)        | kl_target=0.01
+    ------------------------------------------------------------
+    TRPO Relation     | Approximation         | Direct descendant
+    ------------------------------------------------------------
+
+    **When to use PPO-Penalty:**
+    • When you need precise KL control
+    • When you're comfortable tuning kl_target
+    • When you want to stay closer to TRPO theory
+
+    **When to use PPO-Clip:**
+    • Default choice for most problems
+    • Simpler, fewer hyperparameters
+    • More widely adopted in practice
+    """
+
+    print(comparison)
+
+
+if __name__ == "__main__":
+    main()
+    # compare_ppo_variants()  # Uncomment to see comparison

examples/tutorials/rl/cart_pole/step_2_reinforce.py

@@ -0,0 +1,382 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+策略梯度法
+
+
+如果在相同的动作序列下，环境会输出相同的状态和奖励，此方法会彻底失效。
+
+推车立杆任体力比较简单，每次只有2个动作，它很容易就将每一步的最优选择学会了，但对于复杂任务，REINFORCE 可能要困难得多得多。
+
+其本质是搜集多局游戏数据，进行奖励最大化，再依赖环境的随机性进一步迭代。其类似于动态规划算法。 并不保证能找到全局最优解。
+
+
+"""
+import gymnasium as gym
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.distributions import Categorical
+import numpy as np
+import matplotlib.pyplot as plt
+from collections import deque
+import warnings
+
+warnings.filterwarnings('ignore')
+
+
+# ==================== 策略网络 ====================
+class PolicyNetwork(nn.Module):
+    """
+    策略网络：状态 -> 动作概率分布
+    输出是每个动作的概率（经过softmax）
+    """
+
+    def __init__(self, state_dim, hidden_dim, action_dim):
+        super(PolicyNetwork, self).__init__()
+
+        self.network = nn.Sequential(
+            nn.Linear(state_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, action_dim),
+            nn.Softmax(dim=-1)  # 输出概率分布
+        )
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.orthogonal_(module.weight, gain=np.sqrt(2))
+            nn.init.constant_(module.bias, 0)
+
+    def forward(self, state):
+        """
+        输入: state [batch_size, state_dim]
+        输出: action_probs [batch_size, action_dim]
+        """
+        return self.network(state)
+
+    def get_action(self, state):
+        """
+        根据概率分布采样动作
+        返回: action, log_prob
+        """
+        state = torch.FloatTensor(state).unsqueeze(0)
+        # state shape: [1, 4]
+        probs = self.forward(state)
+        # probs shape: [1, 2]
+        dist = Categorical(probs)
+        action = dist.sample()
+        # action shape: [1], 一个数值。
+        log_prob = dist.log_prob(action)
+        # log_prob shape: [1], 一个数值。
+        return action.item(), log_prob
+
+
+# ==================== REINFORCE智能体 ====================
+class ReinforceAgent:
+    """
+    REINFORCE: 蒙特卡洛策略梯度
+    核心思想：用完整轨迹的累积奖励来更新策略
+    好的轨迹 -> 增加这些动作的概率
+    坏的轨迹 -> 减少这些动作的概率
+    """
+
+    def __init__(self,
+                 state_dim,
+                 hidden_dim=128,
+                 action_dim=2,
+                 lr=1e-3,
+                 gamma=0.99):
+
+        self.gamma = gamma
+        self.policy = PolicyNetwork(state_dim, hidden_dim, action_dim)
+        self.optimizer = optim.Adam(self.policy.parameters(), lr=lr)
+
+        # 轨迹存储
+        self.log_probs = []  # 每个动作的对数概率
+        self.rewards = []  # 每个时间步的奖励
+
+    def select_action(self, state):
+        """选择动作并记录对数概率"""
+        action, log_prob = self.policy.get_action(state)
+        self.log_probs.append(log_prob)
+        return action
+
+    def store_reward(self, reward):
+        """存储奖励"""
+        self.rewards.append(reward)
+
+    def update(self):
+        """
+        REINFORCE更新公式：
+        ∇J = E[ Σ ∇logπ(a_t|s_t) * G_t ]
+        其中 G_t = Σ γ^(k-t) * r_k 是从t开始的累积折扣奖励
+        """
+        # 计算累积折扣奖励 G_t
+        returns = []
+        G = 0
+        for r in reversed(self.rewards):
+            G = r + self.gamma * G
+            returns.insert(0, G)
+
+        returns = torch.tensor(returns)
+
+        # 标准化 returns（降低方差，不是必须但很有帮助）
+        returns = (returns - returns.mean()) / (returns.std() + 1e-9)
+
+        # 计算策略梯度损失
+        policy_loss = []
+        for log_prob, G in zip(self.log_probs, returns):
+            # 核心公式：-log_prob * G
+            # 负号是因为PyTorch做梯度下降，我们要最大化J
+            policy_loss.append(-log_prob * G)
+
+        policy_loss = torch.stack(policy_loss).sum()
+
+        # 更新策略
+        self.optimizer.zero_grad()
+        policy_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.policy.parameters(), max_norm=0.5)
+        self.optimizer.step()
+
+        # 清空轨迹
+        self.log_probs = []
+        self.rewards = []
+
+        return policy_loss.item()
+
+    def save(self, path):
+        torch.save(self.policy.state_dict(), path)
+
+    def load(self, path):
+        self.policy.load_state_dict(torch.load(path))
+
+
+# ==================== 训练函数 ====================
+def train_reinforce(env_name='CartPole-v1',
+                    hidden_dim=128,
+                    lr=1e-3,
+                    gamma=0.99,
+                    max_episodes=1000,
+                    log_interval=20):
+    """
+    训练REINFORCE智能体
+    """
+    # 创建环境
+    # env = gym.make(env_name, render_mode='human')
+    env = gym.make(env_name)
+    state_dim = env.observation_space.shape[0]
+    action_dim = env.action_space.n
+
+    # 初始化智能体
+    agent = ReinforceAgent(
+        state_dim=state_dim,
+        hidden_dim=hidden_dim,
+        action_dim=action_dim,
+        lr=lr,
+        gamma=gamma
+    )
+
+    # 记录训练过程
+    episode_rewards = []
+    episode_losses = []
+    moving_avg_rewards = deque(maxlen=100)
+
+    print(f"开始训练 REINFORCE on {env_name}")
+    print(f"状态维度: {state_dim}, 动作维度: {action_dim}")
+    print(f"学习率: {lr}, 折扣因子: {gamma}")
+    print("-" * 50)
+
+    for episode in range(1, max_episodes + 1):
+        state, _ = env.reset()
+        episode_reward = 0
+        episode_loss = 0
+        done = False
+
+        # 收集一条完整轨迹
+        while not done:
+            # 选择动作
+            action = agent.select_action(state)
+
+            # 执行动作
+            next_state, reward, terminated, truncated, _ = env.step(action)
+            done = terminated or truncated
+
+            # 存储奖励
+            agent.store_reward(reward)
+
+            state = next_state
+            episode_reward += reward
+
+        # 一条轨迹结束后更新策略
+        loss = agent.update()
+        episode_loss = loss
+
+        # 记录
+        episode_rewards.append(episode_reward)
+        episode_losses.append(episode_loss)
+        moving_avg_rewards.append(episode_reward)
+
+        # 打印进度
+        if episode % log_interval == 0:
+            avg_reward = np.mean(moving_avg_rewards)
+            print(f"Episode {episode:5d} | "
+                  f"Reward: {episode_reward:6.2f} | "
+                  f"Avg Reward: {avg_reward:6.2f} | "
+                  f"Loss: {episode_loss:8.4f}")
+
+        # 早停：如果连续100局平均分>=475
+        if len(moving_avg_rewards) == 100 and np.mean(moving_avg_rewards) >= 475:
+            print(f"\n🎉 在第 {episode} 回合解决问题！平均奖励: {np.mean(moving_avg_rewards):.2f}")
+            break
+
+    env.close()
+    return agent, episode_rewards, episode_losses
+
+
+# ==================== 可视化函数 ====================
+def plot_training(rewards, losses, save_path=None):
+    """
+    绘制训练曲线
+    """
+    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
+
+    # 1. 原始奖励曲线
+    axes[0, 0].plot(rewards, alpha=0.6, color='blue', linewidth=0.8)
+    axes[0, 0].set_xlabel('Episode')
+    axes[0, 0].set_ylabel('Total Reward')
+    axes[0, 0].set_title('Training Rewards')
+    axes[0, 0].grid(True, alpha=0.3)
+
+    # 2. 移动平均奖励
+    window = 20
+    moving_avg = np.convolve(rewards, np.ones(window) / window, mode='valid')
+    axes[0, 1].plot(moving_avg, color='red', linewidth=2)
+    axes[0, 1].fill_between(range(len(moving_avg)),
+                            moving_avg - np.std(rewards[:len(moving_avg)]),
+                            moving_avg + np.std(rewards[:len(moving_avg)]),
+                            alpha=0.2, color='red')
+    axes[0, 1].set_xlabel('Episode')
+    axes[0, 1].set_ylabel(f'Moving Avg Reward (window={window})')
+    axes[0, 1].set_title('Smoothed Training Curve')
+    axes[0, 1].grid(True, alpha=0.3)
+
+    # 3. 损失曲线
+    axes[1, 0].plot(losses, color='green', alpha=0.6, linewidth=0.8)
+    axes[1, 0].set_xlabel('Episode')
+    axes[1, 0].set_ylabel('Policy Loss')
+    axes[1, 0].set_title('Training Loss')
+    axes[1, 0].grid(True, alpha=0.3)
+
+    # 4. 奖励分布直方图
+    axes[1, 1].hist(rewards[-100:], bins=20, color='purple', alpha=0.7, edgecolor='black')
+    axes[1, 1].set_xlabel('Total Reward')
+    axes[1, 1].set_ylabel('Frequency')
+    axes[1, 1].set_title('Reward Distribution (Last 100 Episodes)')
+    axes[1, 1].axvline(x=np.mean(rewards[-100:]), color='red', linestyle='--',
+                       label=f'Mean: {np.mean(rewards[-100:]):.1f}')
+    axes[1, 1].legend()
+    axes[1, 1].grid(True, alpha=0.3)
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=150, bbox_inches='tight')
+    plt.show()
+
+
+# ==================== 演示智能体 ====================
+def demo_agent(agent, env_name='CartPole-v1', episodes=5):
+    """
+    演示训练好的智能体
+    """
+    env = gym.make(env_name, render_mode='human')
+    # env = gym.make(env_name)
+
+    for episode in range(episodes):
+        state, _ = env.reset()
+        total_reward = 0
+        done = False
+
+        while not done:
+            action, _ = agent.policy.get_action(state)
+            state, reward, terminated, truncated, _ = env.step(action)
+            done = terminated or truncated
+            total_reward += reward
+
+        print(f"Demo Episode {episode + 1}: Reward = {total_reward}")
+
+    env.close()
+
+
+# ==================== 超参数调优 ====================
+def hyperparameter_sweep():
+    """
+    简单的超参数搜索
+    """
+    learning_rates = [1e-4, 3e-4, 1e-3, 3e-3]
+    hidden_sizes = [64, 128, 256]
+
+    results = {}
+
+    for lr in learning_rates:
+        for hidden in hidden_sizes:
+            print(f"\n测试 lr={lr}, hidden={hidden}")
+            _, rewards, _ = train_reinforce(
+                lr=lr,
+                hidden_dim=hidden,
+                max_episodes=300,
+                log_interval=50
+            )
+
+            avg_reward = np.mean(rewards[-50:])
+            results[(lr, hidden)] = avg_reward
+            print(f"平均奖励: {avg_reward:.2f}")
+
+    # 找出最佳参数
+    best_params = max(results, key=results.get)
+    print(f"\n最佳参数: lr={best_params[0]}, hidden={best_params[1]}")
+    print(f"最佳平均奖励: {results[best_params]:.2f}")
+
+    return results
+
+
+# ==================== 主程序 ====================
+if __name__ == "__main__":
+    # 设置随机种子（可复现）
+    torch.manual_seed(42)
+    np.random.seed(42)
+
+    # 训练参数
+    CONFIG = {
+        'env_name': 'CartPole-v1',
+        'hidden_dim': 128,
+        'lr': 1e-3,
+        'gamma': 0.99,
+        'max_episodes': 800,
+        'log_interval': 20
+    }
+
+    # 训练智能体
+    agent, rewards, losses = train_reinforce(**CONFIG)
+
+    # 绘制训练曲线
+    plot_training(rewards, losses, save_path='reinforce_training.png')
+
+    # 打印最终结果
+    print("\n" + "=" * 50)
+    print("训练完成！")
+    print(f"最高奖励: {max(rewards):.2f}")
+    print(f"平均奖励(最后100局): {np.mean(rewards[-100:]):.2f}")
+    print(f"标准差(最后100局): {np.std(rewards[-100:]):.2f}")
+    print("=" * 50)
+
+    # 保存模型
+    agent.save('reinforce_cartpole.pth')
+    print("模型已保存到 reinforce_cartpole.pth")
+
+    # 演示
+    print("\n开始演示...")
+    demo_agent(agent, episodes=3)

examples/tutorials/rl/cart_pole/step_2_reinforce_with_baseline.py

@@ -0,0 +1,332 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+带基线的方法（REINFORCE with Baseline），相当于是先预测一个未来预期奖励的估计，
+然后如果实际的奖励大于这个值，则模型得到正向反馈，动作的概率会被加大，
+如果小于这个值，则动作的概率会减小。
+这相比于原始的方法（REINFORCE）永远只是增大动作的概率使训练变得更稳定。
+
+如果按照训练的轮次/游戏局数来比较，带基线的方法确实收敛更快。
+但额外的价值网络会增加额外的价值网络。
+
+原始REINFORCE 方法，当确定的动作序列导致确定的状态序列时，会失效。
+但带基线REINFORCE 仍然有效。
+
+"""
+import argparse
+from collections import deque
+
+import gymnasium as gym
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.distributions import Categorical
+
+
+class PolicyNetwork(nn.Module):
+    """策略网络 - 输出动作概率分布"""
+
+    def __init__(self, state_dim, hidden_dim, action_dim):
+        super(PolicyNetwork, self).__init__()
+        self.fc1 = nn.Linear(state_dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.fc3 = nn.Linear(hidden_dim, action_dim)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, state):
+        x = torch.relu(self.fc1(state))
+        x = torch.relu(self.fc2(x))
+        x = self.fc3(x)
+        return self.softmax(x)
+
+
+class ValueNetwork(nn.Module):
+    """价值网络 - 作为基线，估计状态价值"""
+
+    def __init__(self, state_dim, hidden_dim):
+        super(ValueNetwork, self).__init__()
+        self.fc1 = nn.Linear(state_dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.fc3 = nn.Linear(hidden_dim, 1)
+
+    def forward(self, state):
+        x = torch.relu(self.fc1(state))
+        x = torch.relu(self.fc2(x))
+        return self.fc3(x)
+
+
+class REINFORCEwithBaseline:
+    """带基线的策略梯度算法"""
+
+    def __init__(self,
+                 env,
+                 policy_lr=1e-3,
+                 value_lr=1e-3,
+                 gamma=0.99,
+                 hidden_dim=128,
+                 render=False):
+
+        self.env = env
+        self.gamma = gamma
+        self.render = render
+
+        # 获取状态和动作维度
+        self.state_dim = env.observation_space.shape[0]
+        self.action_dim = env.action_space.n
+
+        # 初始化策略网络和价值网络
+        self.policy_net = PolicyNetwork(self.state_dim, hidden_dim, self.action_dim)
+        self.value_net = ValueNetwork(self.state_dim, hidden_dim)
+
+        # 优化器
+        self.policy_optimizer = optim.Adam(self.policy_net.parameters(), lr=policy_lr)
+        self.value_optimizer = optim.Adam(self.value_net.parameters(), lr=value_lr)
+
+        # 存储轨迹
+        self.reset_memory()
+
+        # 记录训练信息
+        self.training_stats = {'episode_rewards': [], 'baseline_loss': []}
+
+    def reset_memory(self):
+        """重置存储的记忆"""
+        self.states = []
+        self.actions = []
+        self.rewards = []
+        self.log_probs = []
+
+    def select_action(self, state):
+        """根据当前策略选择动作"""
+        state = torch.FloatTensor(state).unsqueeze(0)
+        probs = self.policy_net(state)
+        m = Categorical(probs)
+        action = m.sample()
+        log_prob = m.log_prob(action)
+
+        # 存储经验
+        self.states.append(state)
+        self.actions.append(action)
+        self.log_probs.append(log_prob)
+
+        return action.item()
+
+    def compute_returns(self):
+        """计算折扣回报"""
+        returns = []
+        R = 0
+        for r in reversed(self.rewards):
+            R = r + self.gamma * R
+            returns.insert(0, R)
+        returns = torch.FloatTensor(returns)
+        # 标准化回报以稳定训练
+        returns = (returns - returns.mean()) / (returns.std() + 1e-9)
+        return returns
+
+    def update(self):
+        """更新策略网络和价值网络"""
+        if len(self.rewards) == 0:
+            return
+
+        # 计算回报和价值估计
+        returns = self.compute_returns()
+        states = torch.cat(self.states)
+
+        # 1. 更新价值网络（基线）
+        value_pred = self.value_net(states).squeeze()
+        value_loss = nn.MSELoss()(value_pred, returns)
+
+        self.value_optimizer.zero_grad()
+        value_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.value_net.parameters(), 0.5)
+        self.value_optimizer.step()
+
+        # 2. 更新策略网络
+        # 重新计算价值估计用于优势函数
+        with torch.no_grad():
+            baselines = self.value_net(states).squeeze()
+
+        # 计算策略梯度
+        policy_loss = []
+        for log_prob, ret, baseline in zip(self.log_probs, returns, baselines):
+            advantage = ret - baseline
+            policy_loss.append(-log_prob * advantage)
+
+        policy_loss = torch.cat(policy_loss).sum()
+
+        self.policy_optimizer.zero_grad()
+        policy_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.policy_net.parameters(), 0.5)
+        self.policy_optimizer.step()
+
+        # 记录价值损失
+        self.training_stats['baseline_loss'].append(value_loss.item())
+
+        # 清空记忆
+        self.reset_memory()
+
+    def train(self, num_episodes, max_steps_per_episode=500):
+        """训练智能体"""
+        episode_rewards = []
+        best_avg_reward = -np.inf
+        reward_window = deque(maxlen=100)
+
+        for episode in range(num_episodes):
+            state, _ = self.env.reset()
+            episode_reward = 0
+
+            for step in range(max_steps_per_episode):
+                if self.render:
+                    self.env.render()
+
+                action = self.select_action(state)
+                next_state, reward, terminated, truncated, _ = self.env.step(action)
+                done = terminated or truncated
+
+                self.rewards.append(reward)
+                episode_reward += reward
+
+                if done:
+                    break
+
+                state = next_state
+
+            # 更新网络
+            self.update()
+
+            # 记录并输出训练信息
+            episode_rewards.append(episode_reward)
+            reward_window.append(episode_reward)
+            avg_reward = np.mean(reward_window)
+
+            self.training_stats['episode_rewards'].append(episode_reward)
+
+            if (episode + 1) % 10 == 0:
+                avg_baseline_loss = np.mean(self.training_stats['baseline_loss'][-10:])
+                print(f'Episode {episode + 1}/{num_episodes}, '
+                      f'Reward: {episode_reward:.1f}, '
+                      f'Avg Reward: {avg_reward:.1f}, '
+                      f'Baseline Loss: {avg_baseline_loss:.4f}')
+
+            # 保存最佳模型
+            if avg_reward > best_avg_reward and episode > 100:
+                best_avg_reward = avg_reward
+                torch.save({
+                    'policy_state_dict': self.policy_net.state_dict(),
+                    'value_state_dict': self.value_net.state_dict(),
+                    'avg_reward': avg_reward,
+                    'episode': episode
+                }, 'best_model.pth')
+
+        return episode_rewards
+
+
+def plot_training_results(rewards, baseline_loss, window=100):
+    """绘制训练结果"""
+    fig, axes = plt.subplots(2, 1, figsize=(10, 8))
+
+    # 绘制奖励曲线
+    axes[0].plot(rewards, alpha=0.3, color='blue', label='Episode Reward')
+
+    # 绘制平滑曲线
+    if len(rewards) >= window:
+        smoothed = np.convolve(rewards, np.ones(window) / window, mode='valid')
+        axes[0].plot(range(window - 1, len(rewards)), smoothed,
+                     color='red', linewidth=2, label=f'Moving Avg (window={window})')
+
+    axes[0].set_xlabel('Episode')
+    axes[0].set_ylabel('Total Reward')
+    axes[0].set_title('REINFORCE with Baseline - Training Rewards')
+    axes[0].legend()
+    axes[0].grid(True, alpha=0.3)
+
+    # 绘制基线损失
+    axes[1].plot(baseline_loss, color='green', alpha=0.6, label='Baseline Loss')
+    axes[1].set_xlabel('Update Step')
+    axes[1].set_ylabel('MSE Loss')
+    axes[1].set_title('Value Network (Baseline) Loss')
+    axes[1].legend()
+    axes[1].grid(True, alpha=0.3)
+
+    plt.tight_layout()
+    plt.savefig('training_results.png', dpi=100)
+    plt.show()
+
+
+def test_agent(env, policy_net, num_episodes=10, render=True):
+    """测试训练好的智能体"""
+    episode_rewards = []
+
+    for episode in range(num_episodes):
+        state, _ = env.reset()
+        episode_reward = 0
+        done = False
+
+        while not done:
+            if render:
+                env.render()
+
+            state_tensor = torch.FloatTensor(state).unsqueeze(0)
+            with torch.no_grad():
+                probs = policy_net(state_tensor)
+                action = torch.argmax(probs).item()
+
+            next_state, reward, terminated, truncated, _ = env.step(action)
+            done = terminated or truncated
+            episode_reward += reward
+            state = next_state
+
+        episode_rewards.append(episode_reward)
+        print(f'Test Episode {episode + 1}: Reward = {episode_reward}')
+
+    print(f'Average Test Reward: {np.mean(episode_rewards):.2f} +/- {np.std(episode_rewards):.2f}')
+    return episode_rewards
+
+
+def main():
+    parser = argparse.ArgumentParser(description='REINFORCE with Baseline for CartPole')
+    parser.add_argument('--episodes', type=int, default=1000, help='Number of training episodes')
+    parser.add_argument('--policy_lr', type=float, default=1e-3, help='Learning rate for policy network')
+    parser.add_argument('--value_lr', type=float, default=1e-3, help='Learning rate for value network')
+    parser.add_argument('--gamma', type=float, default=0.99, help='Discount factor')
+    parser.add_argument('--hidden_dim', type=int, default=128, help='Hidden layer dimension')
+    parser.add_argument('--render_train', action='store_true', help='Render during training')
+    parser.add_argument('--render_test', action='store_true', help='Render during testing')
+    args = parser.parse_args()
+
+    # 创建环境
+    env = gym.make('CartPole-v1')
+
+    # 创建智能体
+    agent = REINFORCEwithBaseline(
+        env=env,
+        policy_lr=args.policy_lr,
+        value_lr=args.value_lr,
+        gamma=args.gamma,
+        hidden_dim=args.hidden_dim,
+        render=args.render_train
+    )
+
+    # 训练
+    print("开始训练 REINFORCE with Baseline...")
+    rewards = agent.train(num_episodes=args.episodes)
+
+    # 绘制训练结果
+    plot_training_results(
+        agent.training_stats['episode_rewards'],
+        agent.training_stats['baseline_loss']
+    )
+
+    # 测试
+    print("\n测试训练好的模型...")
+    test_env = gym.make('CartPole-v1')
+    test_agent(test_env, agent.policy_net, num_episodes=10, render=args.render_test)
+
+    # 关闭环境
+    env.close()
+    test_env.close()
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rl/cart_pole/step_2_rl_dqn.py

@@ -0,0 +1,251 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+深度强化学习 + 基于值 + 模型无关 + 异策略 + 离线学习
+
+强化学习算法
+│
+├── 🔵 基于值的方法 (Value-Based)
+│   ├── 传统: Q-learning, SARSA
+│   └── 🎯 DQN ← 在这里！
+│       ├── DQN
+│       ├── Double DQN
+│       ├── Dueling DQN
+│       └── Rainbow
+│
+├── 🔴 基于策略的方法 (Policy-Based)
+│   ├── REINFORCE
+│   ├── PPO
+│   └── TRPO
+│
+└── 🟣 演员-评论家方法 (Actor-Critic)
+    ├── A2C/A3C
+    ├── SAC
+    └── TD3
+
+贝尔曼方程
+
+贝尔曼的洞察（1957）：
+"最优策略有这样的性质：无论初始状态和初始决策如何，其余的决策必须构成一个以第一个决策产生的状态为初始状态的最优策略。"
+翻译成人话：
+如果A→B→C是最优路径，那么B→C也必须是最优路径！
+
+备注：所以 DQN 是基于贝尔曼假设的，当环境不符合这个假设时，该方法不成立。
+
+"""
+from collections import deque
+import random
+
+import gymnasium as gym
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+
+# 神经网络定义
+class DQN(nn.Module):
+    def __init__(self, state_size: int, action_size: int):
+        super(DQN, self).__init__()
+        self.fc1 = nn.Linear(state_size, 64)
+        self.fc2 = nn.Linear(64, 64)
+        self.fc3 = nn.Linear(64, action_size)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x = self.relu(self.fc1(x))
+        x = self.relu(self.fc2(x))
+        return self.fc3(x)
+
+
+# 经验回放缓冲区
+class ReplayBuffer:
+    def __init__(self, capacity: int):
+        self.buffer = deque(maxlen=capacity)
+
+    def push(self, state, action, reward, next_state, done):
+        """
+        :param state: 状态 = [位置, 速度, 角度, 角速度]
+        :param action:
+        :param reward: float
+        :param next_state:
+        :param done:
+        :return:
+        """
+        self.buffer.append((state, action, reward, next_state, done))
+
+    def sample(self, batch_size):
+        batch = random.sample(self.buffer, batch_size)
+        states, actions, rewards, next_states, dones = zip(*batch)
+        result = (
+            np.array(states),
+            np.array(actions),
+            np.array(rewards),
+            np.array(next_states),
+            np.array(dones)
+        )
+        return result
+
+    def __len__(self):
+        return len(self.buffer)
+
+
+# DQN智能体
+class DQNAgent:
+    def __init__(self, state_size: int, action_size: int):
+        self.state_size = state_size
+        self.action_size = action_size
+
+        # 网络
+        self.policy_net = DQN(state_size, action_size)
+        self.target_net = DQN(state_size, action_size)
+        self.target_net.load_state_dict(self.policy_net.state_dict())
+
+        # 优化器
+        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=0.001)
+
+        # 超参数
+        self.gamma = 0.99
+        self.epsilon = 1.0
+        self.epsilon_min = 0.01
+        self.epsilon_decay = 0.995
+        self.batch_size = 64
+        self.buffer = ReplayBuffer(10000)
+
+        # 目标网络更新频率
+        self.target_update = 10
+
+    def select_action(self, state, training=True):
+        if training and np.random.random() < self.epsilon:
+            return random.randrange(self.action_size)
+
+        state = torch.FloatTensor(state).unsqueeze(0)
+        with torch.no_grad():
+            q_values = self.policy_net(state)
+        return q_values.argmax().item()
+
+    def train_step(self):
+        if len(self.buffer) < self.batch_size:
+            return
+
+        # 从缓冲区采样
+        states, actions, rewards, next_states, dones = self.buffer.sample(self.batch_size)
+
+        # 转换为张量
+        states = torch.FloatTensor(states)
+        actions = torch.LongTensor(actions).unsqueeze(1)
+        rewards = torch.FloatTensor(rewards).unsqueeze(1)
+        next_states = torch.FloatTensor(next_states)
+        dones = torch.FloatTensor(dones).unsqueeze(1)
+
+        # 计算当前Q值
+        actions_logits = self.policy_net(states)
+        # actions_logits shape: [batch_size, action_size]
+        current_q = torch.gather(actions_logits, dim=1, index=actions)
+        # current_q shape: [batch_size, 1]
+
+        # 计算目标Q值
+        with torch.no_grad():
+            next_actions_logits = self.target_net(next_states)
+            # next_actions_logits shape: [batch_size, action_size]
+            next_q, _ = torch.max(next_actions_logits, 1)
+            next_q = torch.unsqueeze(next_q, 1)
+            # next_q shape: [batch_size, 1]
+            # 贝尔曼方程
+            target_q = rewards + (1 - dones) * self.gamma * next_q
+
+        # 计算损失
+        # current_q 预测采取当前动作后未来能获取的总奖励。并对远期奖励降权。
+        # target_q 上一次预测未来能获取的总奖励 = 当前已获得奖励 + 当前预测未来能获取的总奖励
+        loss = nn.MSELoss()(current_q, target_q)
+
+        # 引入 target_net 而不是直接使用 policy_net 是因为：
+        # 每一个 train_step 只训练了一个 batch，但训练本身具有随机性，
+        # 因此可以认为，只在训练多个 batch 之后对模型的优化才真实有效。
+        # 所以设置为定期将 target_net 与 policy_net 同步。
+
+        # 优化
+        self.optimizer.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.policy_net.parameters(), 1.0)
+        self.optimizer.step()
+
+        # 衰减探索率
+        if self.epsilon > self.epsilon_min:
+            self.epsilon *= self.epsilon_decay
+
+
+# 训练函数
+def train_dqn():
+    env = gym.make("CartPole-v1")
+    state_size = env.observation_space.shape[0]
+    action_size = env.action_space.n
+
+    agent = DQNAgent(state_size, action_size)
+    episodes = 500
+    rewards = []
+
+    for episode in range(episodes):
+        state, _ = env.reset()
+        total_reward = 0
+        done = False
+
+        while not done:
+            # 选择动作
+            action = agent.select_action(state)
+
+            # 执行动作
+            next_state, reward, terminated, truncated, _ = env.step(action)
+            done = terminated or truncated
+
+            # 存储经验
+            agent.buffer.push(state, action, reward, next_state, done)
+
+            # 训练
+            agent.train_step()
+
+            state = next_state
+            total_reward += reward
+
+        # 更新目标网络
+        if episode % agent.target_update == 0:
+            agent.target_net.load_state_dict(agent.policy_net.state_dict())
+
+        rewards.append(total_reward)
+
+        if (episode + 1) % 50 == 0:
+            avg_reward = np.mean(rewards[-50:])
+            print(f"Episode {episode + 1}, Avg Reward: {avg_reward:.2f}, Epsilon: {agent.epsilon:.3f}")
+
+    env.close()
+
+    # 可视化
+    plt.figure(figsize=(12, 4))
+
+    plt.subplot(1, 2, 1)
+    plt.plot(rewards)
+    plt.xlabel('Episode')
+    plt.ylabel('Reward')
+    plt.title('Training Rewards')
+
+    plt.subplot(1, 2, 2)
+    window = 20
+    moving_avg = [np.mean(rewards[max(0, i - window):i + 1]) for i in range(len(rewards))]
+    plt.plot(moving_avg)
+    plt.xlabel('Episode')
+    plt.ylabel('Moving Avg Reward')
+    plt.title(f'Moving Average (window={window})')
+
+    plt.tight_layout()
+    plt.show()
+
+    return agent
+
+
+# 运行训练
+agent = train_dqn()
+
+
+if __name__ == "__main__":
+    pass

examples/tutorials/rlhf/gpt2_sst2/step_1_prepare_data.py

@@ -0,0 +1,59 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+或使用命令行
+pip install modelscope
+modelscope download \
+--model 'qgyd2021/Qwen3-8B-sft-deepspeed' \
+--local_dir '/root/autodl-tmp/trained_models/Qwen3-8B-sft-deepspeed'
+
+"""
+import argparse
+import os
+from pathlib import Path
+import platform
+
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from modelscope import snapshot_download
+# from huggingface_hub import snapshot_download
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo_id", default="Qwen/Qwen2.5-0.5B", type=str)
+    parser.add_argument(
+        "--local_dir",
+        default=(temp_directory / "../trained_models/Qwen/Qwen2.5-0.5B").as_posix(),
+        type=str
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    #modelscope
+    # snapshot_download(
+    #     model_id=args.repo_id,
+    #     local_dir=args.local_dir,
+    # )
+    #huggingface_hub
+    snapshot_download(
+        repo_type="model",
+        repo_id=args.repo_id,
+        local_dir=args.local_dir,
+    )
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2/step_2_train_sft_model.py

@@ -0,0 +1,166 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+用sst的句子训练gpt2模型，让其随机生成一些评论。
+"""
+import argparse
+import os
+from pathlib import Path
+import platform
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from datasets import load_dataset
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, DataCollatorForLanguageModeling, Trainer, TrainingArguments
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_name",
+        # default="openai-community/gpt2",
+        default=(project_path / "pretrained_models/openai-community/gpt2").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--dataset_path",
+        default="stanfordnlp/sst2",
+        type=str
+    ),
+    parser.add_argument("--dataset_name", default=None, type=str),
+    parser.add_argument("--dataset_split", default=None, type=str),
+    parser.add_argument(
+        "--dataset_cache_dir",
+        default=(temp_directory / "hub_datasets").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--model_cache_dir",
+        default=(temp_directory / "hub_models").as_posix(),
+        type=str
+    ),
+    parser.add_argument("--dataset_streaming", default=None, type=str),
+    parser.add_argument("--valid_dataset_size", default=1000, type=int),
+    parser.add_argument("--shuffle_buffer_size", default=5000, type=int),
+
+    parser.add_argument(
+        "--output_model_dir",
+        default=(project_path / "trained_models/gpt2-sst2-generation-20260213-2048").as_posix(),
+        type=str
+    ),
+
+    parser.add_argument(
+        "--num_workers",
+        default=None if platform.system() in ("Windows", "Darwin") else os.cpu_count() // 2,
+        type=int
+    ),
+    parser.add_argument(
+        "--device",
+        default=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        type=int
+    ),
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    model = AutoModelForCausalLM.from_pretrained(args.model_name)
+    model = model.to(args.device)
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    dataset_dict = load_dataset(
+        path=args.dataset_path,
+        name=args.dataset_name,
+        split=args.dataset_split,
+        cache_dir=args.dataset_cache_dir,
+        # num_proc=args.num_workers if not args.dataset_streaming else None,
+        streaming=args.dataset_streaming,
+    )
+    train_dataset = dataset_dict["train"]
+    valid_dataset = dataset_dict["validation"]
+    # test_dataset = dataset_dict["test"]
+
+    def format_func(example):
+        sentence = example["sentence"]
+        sentence += tokenizer.eos_token
+        tokenized = tokenizer(sentence)
+        input_ids = tokenized["input_ids"]
+        attention_mask = tokenized["attention_mask"]
+        # print(input_ids)
+        # print(attention_mask)
+        result = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+        }
+        return result
+
+    train_dataset = train_dataset.map(
+        format_func,
+        batched=False,
+        remove_columns=train_dataset.column_names,
+    )
+    valid_dataset = valid_dataset.map(
+        format_func,
+        batched=False,
+        remove_columns=valid_dataset.column_names,
+    )
+    train_dataset = train_dataset.filter(
+        function=lambda x: len(x["input_ids"]) > 5
+    )
+    valid_dataset = valid_dataset.filter(
+        function=lambda x: len(x["input_ids"]) > 5
+    )
+
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer,
+        mlm=False
+    )
+
+    training_args = TrainingArguments(
+        output_dir=args.output_model_dir,
+        # overwrite_output_dir=True,
+        num_train_epochs=1,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        eval_strategy="steps",
+        eval_steps=100,
+        save_strategy="steps",
+        save_steps=100,
+        save_total_limit=2,
+        logging_steps=100,
+        learning_rate=5e-5,
+        warmup_steps=500,
+        weight_decay=0.01,
+        fp16=torch.cuda.is_available(),
+        dataloader_num_workers=args.num_workers or 0,
+        remove_unused_columns=False,
+        load_best_model_at_end=False,
+        # metric_for_best_model="eval_loss",
+        # greater_is_better=False,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        data_collator=data_collator,
+        train_dataset=train_dataset,
+        eval_dataset=valid_dataset,
+        tokenizer=tokenizer,
+    )
+
+    trainer.train()
+    trainer.save_model()
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2/step_3_train_reward_model.py

@@ -0,0 +1,295 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import os
+from pathlib import Path
+import platform
+from typing import Any, Dict, List, Optional, Union, Tuple
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from datasets import load_dataset
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import (AutoModelForCausalLM,
+                          AutoTokenizer,
+                          GPT2PreTrainedModel, GPT2Config, GPT2Model,
+                          DataCollatorWithPadding,
+                          Trainer, TrainingArguments
+                          )
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_name",
+        # default="openai-community/gpt2",
+        default=(project_path / "pretrained_models/openai-community/gpt2").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--dataset_path",
+        default="stanfordnlp/sst2",
+        type=str
+    ),
+    parser.add_argument("--dataset_name", default=None, type=str),
+    parser.add_argument("--dataset_split", default=None, type=str),
+    parser.add_argument(
+        "--dataset_cache_dir",
+        default=(temp_directory / "hub_datasets").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--model_cache_dir",
+        default=(temp_directory / "hub_models").as_posix(),
+        type=str
+    ),
+    parser.add_argument("--dataset_streaming", default=None, type=str),
+    parser.add_argument("--valid_dataset_size", default=1000, type=int),
+    parser.add_argument("--shuffle_buffer_size", default=5000, type=int),
+
+    parser.add_argument(
+        "--output_model_dir",
+        default=(project_path / "trained_models/gpt2-sst2-reward-20260213-2122").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--num_workers",
+        default=None if platform.system() in ("Windows", "Darwin") else os.cpu_count() // 2,
+        type=int
+    ),
+    parser.add_argument(
+        "--device",
+        default=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        type=str
+    ),
+    args = parser.parse_args()
+    return args
+
+
+class RewardHead(nn.Module):
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.linear = nn.Linear(self.hidden_size, 1)
+        self._post_init()
+
+    def _post_init(self):
+        nn.init.normal_(
+            self.linear.weight,
+            std=(1.0 / np.sqrt(self.hidden_size + 1))
+        )
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(self, hidden_states):
+        # hidden_states shape: [batch_size, seq_len, hidden_size]
+        reward_logits = self.linear(hidden_states)
+        # reward_logits shape: [batch_size, seq_len, 1]
+        return reward_logits
+
+
+class GPT2RewardModel(GPT2PreTrainedModel):
+    def __init__(self, config: GPT2Config):
+        super().__init__(config)
+        self.transformer = GPT2Model(config)
+        self.reward_head = RewardHead(config.hidden_size)
+        self.post_init()
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> Union[Tuple, torch.Tensor]:
+        transformer_outputs = self.transformer(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True
+        )
+        last_hidden_state = transformer_outputs.hidden_states[-1]
+        # last_hidden_state shape: [batch_size, seq_len, hidden_size]
+        rewards_logits = self.reward_head(last_hidden_state)
+        # rewards_logits shape: [batch_size, seq_len, 1]
+        rewards_logits = torch.squeeze(rewards_logits, -1)
+        # rewards_logits shape: [batch_size, seq_len]
+        rewards = torch.sigmoid(rewards_logits)
+        # rewards shape: [batch_size, seq_len]
+        return rewards
+
+    @classmethod
+    def from_pretrained(cls, model_name_or_path, *model_args, **kwargs):
+        config = GPT2Config.from_pretrained(model_name_or_path)
+        model = cls(config)
+        pretrained_model = GPT2Model.from_pretrained(model_name_or_path)
+        model.transformer.load_state_dict(pretrained_model.state_dict(), strict=False)
+        return model
+
+
+class SST2RewardTrainer(Trainer):
+    def compute_loss(
+        self,
+        model: nn.Module,
+        inputs: dict[str, Union[torch.Tensor, Any]],
+        return_outputs: bool = False,
+        num_items_in_batch: Optional[torch.Tensor] = None,
+    ):
+        rewards = model(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"]
+        )
+        sequence_lengths = inputs["attention_mask"].sum(dim=1) - 1
+        batch_indices = torch.arange(rewards.size(0), device=rewards.device)
+        sequence_reward = rewards[batch_indices, sequence_lengths]
+        # sequence_reward shape: [batch_size,]
+        loss = F.mse_loss(
+            sequence_reward,
+            inputs["score"].float()
+        )
+        if return_outputs:
+            return loss, {
+                "loss": loss,
+                "predictions": sequence_reward.detach(),
+            }
+        return loss
+
+    def prediction_step(
+        self,
+        model: nn.Module,
+        inputs: dict[str, Union[torch.Tensor, Any]],
+        prediction_loss_only: bool,
+        ignore_keys: Optional[list[str]] = None,
+    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        with torch.no_grad():
+            loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
+
+        if prediction_loss_only:
+            return loss, None, None
+
+        predictions = outputs["predictions"]
+        labels = inputs["score"].float()
+        return loss, predictions, labels
+
+
+def compute_metrics(eval_pred):
+    """计算评估指标"""
+    predictions, labels = eval_pred
+    predictions = torch.tensor(predictions)
+    labels = torch.tensor(labels)
+
+    error = (predictions - labels).abs()
+
+    return {
+        "mean_error": error.mean().item(),
+        "std_error": error.std().item(),
+        "reward_mean": predictions.mean().item(),
+        "reward_min": predictions.min().item(),
+        "reward_max": predictions.max().item(),
+        "score_mean": labels.mean().item(),
+    }
+
+
+def main():
+    args = get_args()
+
+    model = GPT2RewardModel.from_pretrained(args.model_name)
+    model = model.to(args.device)
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    # tokenized = tokenizer(
+    #     "this is a good day",
+    #     # "this is ",
+    #     return_tensors="pt"
+    # )
+    # output_dict = model(**tokenized)
+
+    dataset_dict = load_dataset(
+        path=args.dataset_path,
+        name=args.dataset_name,
+        split=args.dataset_split,
+        cache_dir=args.dataset_cache_dir,
+        # num_proc=args.num_workers if not args.dataset_streaming else None,
+        streaming=args.dataset_streaming,
+    )
+    train_dataset = dataset_dict["train"]
+    valid_dataset = dataset_dict["validation"]
+    # test_dataset = dataset_dict["test"]
+
+    def format_func(example):
+        sentence: str = example["sentence"]
+        score: float = float(example["label"])
+        sentence += tokenizer.eos_token
+        tokenized = tokenizer(sentence)
+        input_ids = tokenized["input_ids"]
+        attention_mask = tokenized["attention_mask"]
+        result = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "score": score,
+        }
+        return result
+
+    train_dataset = train_dataset.map(
+        format_func,
+        batched=False,
+        remove_columns=train_dataset.column_names,
+    )
+    valid_dataset = valid_dataset.map(
+        format_func,
+        batched=False,
+        remove_columns=valid_dataset.column_names,
+    )
+    train_dataset = train_dataset.filter(
+        function=lambda x: len(x["input_ids"]) > 6
+    )
+    valid_dataset = valid_dataset.filter(
+        function=lambda x: len(x["input_ids"]) > 6
+    )
+    data_collator = DataCollatorWithPadding(tokenizer)
+
+    training_args = TrainingArguments(
+        output_dir=args.output_model_dir,
+        # overwrite_output_dir=True,
+        num_train_epochs=1,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        eval_strategy="steps",
+        eval_steps=500,
+        save_strategy="steps",
+        save_steps=500,
+        save_total_limit=2,
+        logging_steps=500,
+        learning_rate=5e-5,
+        warmup_steps=1000,
+        weight_decay=0.01,
+        fp16=torch.cuda.is_available(),
+        dataloader_num_workers=args.num_workers or 0,
+        remove_unused_columns=False,
+        load_best_model_at_end=True,
+        metric_for_best_model="eval_loss",
+        greater_is_better=False,
+    )
+
+    trainer = SST2RewardTrainer(
+        model=model,
+        args=training_args,
+        data_collator=data_collator,
+        train_dataset=train_dataset,
+        eval_dataset=valid_dataset,
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics,
+    )
+
+    trainer.train()
+    trainer.save_model()
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2/step_4_test_reward_model.py

@@ -0,0 +1,160 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import os
+from pathlib import Path
+import platform
+from typing import Any, Dict, List, Optional, Union, Tuple
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from datasets import load_dataset
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import (AutoTokenizer,
+                          GPT2PreTrainedModel, GPT2Config, GPT2Model,
+                          )
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_name",
+        # default="openai-community/gpt2",
+        # default=(project_path / "trained_models/gpt2-sst2-reward").as_posix(),
+        default=(project_path / "trained_models/gpt2-sst2-reward-20260213-2122").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--dataset_path",
+        default="stanfordnlp/sst2",
+        type=str
+    ),
+    parser.add_argument("--dataset_name", default=None, type=str),
+    parser.add_argument("--dataset_split", default=None, type=str),
+    parser.add_argument(
+        "--dataset_cache_dir",
+        default=(temp_directory / "hub_datasets").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--model_cache_dir",
+        default=(temp_directory / "hub_models").as_posix(),
+        type=str
+    ),
+    parser.add_argument("--dataset_streaming", default=None, type=str),
+
+    parser.add_argument(
+        "--num_workers",
+        default=None if platform.system() in ("Windows", "Darwin") else os.cpu_count() // 2,
+        type=int
+    ),
+    parser.add_argument(
+        "--device",
+        default=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        type=str
+    ),
+    args = parser.parse_args()
+    return args
+
+
+class RewardHead(nn.Module):
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.linear = nn.Linear(self.hidden_size, 1)
+        self._post_init()
+
+    def _post_init(self):
+        nn.init.normal_(
+            self.linear.weight,
+            std=(1.0 / np.sqrt(self.hidden_size + 1))
+        )
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(self, hidden_states):
+        # hidden_states shape: [batch_size, seq_len, hidden_size]
+        reward_logits = self.linear(hidden_states)
+        # reward_logits shape: [batch_size, seq_len, 1]
+        return reward_logits
+
+
+class GPT2RewardModel(GPT2PreTrainedModel):
+    def __init__(self, config: GPT2Config):
+        super().__init__(config)
+        self.transformer = GPT2Model(config)
+        self.reward_head = RewardHead(config.hidden_size)
+        self.post_init()
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> Union[Tuple, torch.Tensor]:
+        transformer_outputs = self.transformer(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True
+        )
+        last_hidden_state = transformer_outputs.hidden_states[-1]
+        # last_hidden_state shape: [batch_size, seq_len, hidden_size]
+        rewards_logits = self.reward_head(last_hidden_state)
+        # rewards_logits shape: [batch_size, seq_len, 1]
+        rewards_logits = torch.squeeze(rewards_logits, -1)
+        # rewards_logits shape: [batch_size, seq_len]
+        rewards = torch.sigmoid(rewards_logits)
+        # rewards shape: [batch_size, seq_len]
+        return rewards
+
+
+def main():
+    args = get_args()
+
+    model = GPT2RewardModel.from_pretrained(args.model_name)
+    model = model.to(args.device)
+    # model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    dataset_dict = load_dataset(
+        path=args.dataset_path,
+        name=args.dataset_name,
+        split=args.dataset_split,
+        cache_dir=args.dataset_cache_dir,
+        # num_proc=args.num_workers if not args.dataset_streaming else None,
+        streaming=args.dataset_streaming,
+    )
+    # dataset = dataset_dict["train"]
+    dataset = dataset_dict["validation"]
+    # dataset = dataset_dict["test"]
+
+    for example in dataset:
+        sentence: str = example["sentence"]
+        score: float = float(example["label"])
+        # sentence = "this is very good movie, I recommend it."
+
+        sentence += tokenizer.eos_token
+        tokenized = tokenizer(
+            sentence,
+            return_tensors="pt"
+        )
+        with torch.no_grad():
+            rewards = model(**tokenized)
+        rewards = rewards[0]
+        rewards = rewards.detach().cpu().numpy()
+        last_token_reward = rewards[-1]
+        #rewards: {rewards}\n
+        msg = f"last_token_reward: {last_token_reward}\nscore: {score}\nsentence: {sentence}\n"
+        print(msg)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2/step_5_ppo_rlhf.py

@@ -0,0 +1,430 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import copy
+import os
+import random
+from pathlib import Path
+import platform
+from typing import Optional, Tuple, List, Dict, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer, AutoModelForCausalLM, GPT2PreTrainedModel,
+    GPT2Config, GPT2Model, GPT2LMHeadModel, DataCollatorWithPadding
+)
+
+# 路径配置
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = Path(os.path.abspath("../../../"))
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--reward_model_name", type=str,
+                        default=(project_path / "trained_models/gpt2-sst2-reward").as_posix())
+    parser.add_argument("--sft_model_name", type=str,
+                        default=(project_path / "trained_models/gpt2-sst2-generation").as_posix())
+    parser.add_argument("--dataset_path", default="stanfordnlp/sst2", type=str)
+    parser.add_argument("--dataset_cache_dir",
+                        default=(temp_directory / "hub_datasets").as_posix(), type=str)
+    parser.add_argument("--model_cache_dir",
+                        default=(temp_directory / "hub_models").as_posix(), type=str)
+    parser.add_argument("--valid_dataset_size", default=1000, type=int)
+
+    # 训练参数
+    parser.add_argument("--batch_size", default=16, type=int)  # CPU上用小一点的batch
+    parser.add_argument("--ppo_epochs", default=4, type=int)
+    parser.add_argument("--mini_batch_size", default=4, type=int)
+    parser.add_argument("--kl_beta", default=0.2, type=float)
+    parser.add_argument("--gamma", default=1.0, type=float)
+    parser.add_argument("--lam", default=0.95, type=float)
+    parser.add_argument("--clip_epsilon", default=0.2, type=float)
+    parser.add_argument("--lr", default=1e-5, type=float)
+    parser.add_argument("--max_epochs", default=10, type=int)
+
+    # 生成参数
+    parser.add_argument("--max_new_tokens", default=32, type=int)
+    parser.add_argument("--top_p", default=0.85, type=float)
+    parser.add_argument("--temperature", default=0.85, type=float)
+    parser.add_argument("--min_response_len", default=5, type=int)
+    parser.add_argument("--max_response_len", default=16, type=int)
+
+    # 其他
+    parser.add_argument("--num_workers", default=0 if platform.system() == "Windows" else 2, type=int)
+    parser.add_argument("--device", default="cpu", type=str)  # 强制用CPU
+
+    return parser.parse_args()
+
+
+class ValueHead(nn.Module):
+    """价值头，为每个token预测一个价值"""
+
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.linear = nn.Linear(hidden_size, 1)
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.normal_(self.linear.weight, std=1.0 / np.sqrt(self.linear.in_features + 1))
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(self, hidden_states):
+        return self.linear(hidden_states).squeeze(-1)
+
+
+class GPT2ActorCritic(GPT2PreTrainedModel):
+    """Actor-Critic模型，同时输出logits和values"""
+
+    def __init__(self, config: GPT2Config):
+        super().__init__(config)
+        self.lm = GPT2LMHeadModel(config)
+        self.value_head = ValueHead(config.hidden_size)
+        self.post_init()
+
+    def forward(self, input_ids, attention_mask=None):
+        outputs = self.lm(
+            input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True
+        )
+        # values来自最后一层hidden states
+        values = self.value_head(outputs.hidden_states[-1])
+        return outputs.logits, values
+
+    def generate(self, *args, **kwargs):
+        return self.lm.generate(*args, **kwargs)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name):
+        """从预训练GPT2LMHeadModel加载"""
+        config = GPT2Config.from_pretrained(pretrained_model_name)
+        model = cls(config)
+        pretrained = GPT2LMHeadModel.from_pretrained(pretrained_model_name)
+        model.lm.load_state_dict(pretrained.state_dict(), strict=False)
+        return model
+
+
+class GPT2RewardModel(GPT2PreTrainedModel):
+    """奖励模型，为每个token预测奖励"""
+
+    def __init__(self, config: GPT2Config):
+        super().__init__(config)
+        self.transformer = GPT2Model(config)
+        self.reward_head = nn.Linear(config.hidden_size, 1)
+        self.post_init()
+
+    def forward(self, input_ids, attention_mask=None):
+        outputs = self.transformer(
+            input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True
+        )
+        rewards = self.reward_head(outputs.hidden_states[-1]).squeeze(-1)
+        return torch.sigmoid(rewards)  # [batch, seq_len]
+
+
+class PPOAgent:
+    """PPO训练Agent，封装所有训练逻辑"""
+
+    def __init__(self, args):
+        self.args = args
+        self.device = torch.device(args.device)
+
+        # 加载tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(args.sft_model_name)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
+        # 加载模型
+        print("Loading models...")
+        self.actor_critic = GPT2ActorCritic.from_pretrained(args.sft_model_name).to(self.device)
+        self.reward_model = GPT2RewardModel.from_pretrained(args.reward_model_name).to(self.device)
+        self.reward_model.eval()
+
+        # 参考模型（冻结）
+        self.ref_model = copy.deepcopy(self.actor_critic).to(self.device)
+        self.ref_model.eval()
+
+        # 优化器
+        self.optimizer = torch.optim.Adam(self.actor_critic.parameters(), lr=args.lr)
+
+        # 训练状态
+        self.training_step = 0
+
+    def prepare_dataset(self):
+        """准备训练数据集"""
+        print("Loading dataset...")
+        dataset = load_dataset(
+            path=self.args.dataset_path,
+            cache_dir=self.args.dataset_cache_dir,
+            split="train"
+        )
+
+        def filter_and_truncate(example):
+            # 只保留足够长的句子
+            tokens = self.tokenizer(example["sentence"])["input_ids"]
+            if len(tokens) <= 8:
+                return False
+
+            # 随机截取前2-6个token作为query
+            example["query_ids"] = tokens[:random.randint(2, 6)]
+            return True
+
+        dataset = dataset.filter(filter_and_truncate)
+        dataset = dataset.select(range(min(len(dataset), 5000)))  # CPU上用小数据集
+
+        return dataset
+
+    def collect_rollouts(self, batch):
+        """收集一轮交互数据"""
+        query_ids_list = []
+        response_ids_list = []
+        rewards_list = []
+
+        for i in range(len(batch["query_ids"])):
+            query_ids = torch.tensor(batch["query_ids"][i]).to(self.device)
+            query_ids_list.append(query_ids)
+
+            # 生成response
+            with torch.no_grad():
+                response_len = random.randint(
+                    self.args.min_response_len,
+                    self.args.max_response_len
+                )
+                full_ids = self.actor_critic.generate(
+                    input_ids=query_ids.unsqueeze(0),
+                    max_new_tokens=response_len,
+                    do_sample=True,
+                    top_p=self.args.top_p,
+                    temperature=self.args.temperature,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                )[0]
+
+                response_ids = full_ids[len(query_ids):]
+                response_ids_list.append(response_ids)
+
+                # 计算奖励（只取最后一个token的奖励）
+                reward = self.reward_model(
+                    full_ids.unsqueeze(0),
+                    attention_mask=torch.ones_like(full_ids).unsqueeze(0)
+                )[0, -1]
+                # 缩放到[-1, 1]
+                rewards_list.append(2 * (reward - 0.5))
+
+        return query_ids_list, response_ids_list, rewards_list
+
+    def compute_advantages_and_returns(self, log_probs, values, rewards, masks):
+        """计算GAE advantages和returns"""
+        seq_len = rewards.shape[1]
+        advantages = torch.zeros_like(rewards)
+        returns = torch.zeros_like(rewards)
+
+        gae = 0
+        for t in reversed(range(seq_len)):
+            if t == seq_len - 1:
+                next_value = 0
+            else:
+                next_value = values[:, t + 1]
+
+            delta = rewards[:, t] + self.args.gamma * next_value - values[:, t]
+            gae = delta + self.args.gamma * self.args.lam * gae
+            advantages[:, t] = gae
+            returns[:, t] = advantages[:, t] + values[:, t]
+
+        # 只对有效位置进行whiten
+        advantages = self.masked_whiten(advantages, masks)
+        return advantages, returns
+
+    def masked_whiten(self, values, mask):
+        """带mask的whitening"""
+        mask = mask.float()
+        mean = (values * mask).sum() / mask.sum()
+        var = (((values - mean) * mask) ** 2).sum() / mask.sum()
+        whitened = (values - mean) * torch.rsqrt(var + 1e-8)
+        return whitened * mask
+
+    def ppo_step(self, batch_data):
+        """单步PPO更新"""
+        (query_ids_list, response_ids_list, old_log_probs,
+         advantages, returns, masks) = batch_data
+
+        # 拼接完整的query+response
+        full_ids_list = []
+        for q, r in zip(query_ids_list, response_ids_list):
+            full_ids_list.append(torch.cat([q, r]))
+
+        # padding
+        padded = self.tokenizer.pad(
+            {"input_ids": full_ids_list},
+            padding=True,
+            return_tensors="pt"
+        )
+        input_ids = padded["input_ids"].to(self.device)
+        attention_mask = padded["attention_mask"].to(self.device)
+
+        # 前向传播
+        logits, values = self.actor_critic(input_ids, attention_mask)
+
+        # 计算新的log_probs
+        log_probs = F.log_softmax(logits[:, :-1, :], dim=-1)
+        log_probs = torch.gather(
+            log_probs, 2,
+            input_ids[:, 1:].unsqueeze(-1)
+        ).squeeze(-1)
+
+        # 只保留response部分的log_probs
+        response_start = [len(q) for q in query_ids_list]
+        new_log_probs = []
+        for i, start in enumerate(response_start):
+            new_log_probs.append(log_probs[i, start - 1:start - 1 + len(response_ids_list[i])])
+        new_log_probs = torch.cat(new_log_probs)
+
+        # 计算ratio和PPO损失
+        old_log_probs = old_log_probs.detach()
+        ratio = torch.exp(new_log_probs - old_log_probs)
+
+        # 裁剪的policy loss
+        surr1 = ratio * advantages
+        surr2 = torch.clamp(ratio, 1 - self.args.clip_epsilon,
+                            1 + self.args.clip_epsilon) * advantages
+        policy_loss = -torch.min(surr1, surr2).mean()
+
+        # value loss
+        value_pred = []
+        for i, start in enumerate(response_start):
+            value_pred.append(values[i, start - 1:start - 1 + len(response_ids_list[i])])
+        value_pred = torch.cat(value_pred)
+        value_loss = F.mse_loss(value_pred, returns)
+
+        # 总loss
+        loss = policy_loss + 0.5 * value_loss
+
+        return loss, policy_loss, value_loss
+
+    def train_epoch(self, dataset):
+        """训练一个epoch"""
+        total_policy_loss = 0
+        total_value_loss = 0
+        num_batches = 0
+
+        for batch_idx in range(0, len(dataset), self.args.batch_size):
+            # 1. 收集数据
+            batch = dataset[batch_idx:batch_idx + self.args.batch_size]
+            query_ids_list, response_ids_list, rewards_list = self.collect_rollouts(batch)
+
+            # 2. 计算旧的log_probs和values
+            old_log_probs_list = []
+            values_list = []
+            masks_list = []
+
+            with torch.no_grad():
+                for q_ids, r_ids in zip(query_ids_list, response_ids_list):
+                    full_ids = torch.cat([q_ids, r_ids]).unsqueeze(0).to(self.device)
+                    attn_mask = torch.ones_like(full_ids)
+
+                    logits, values = self.actor_critic(full_ids, attn_mask)
+
+                    # 计算response部分的log_probs
+                    log_probs = F.log_softmax(logits[:, :-1, :], dim=-1)
+                    log_probs = torch.gather(
+                        log_probs, 2,
+                        full_ids[:, 1:].unsqueeze(-1)
+                    ).squeeze(-1)
+
+                    start = len(q_ids) - 1
+                    end = start + len(r_ids)
+                    old_log_probs_list.append(log_probs[0, start:end])
+                    values_list.append(values[0, start:end])
+
+                    # 创建mask
+                    mask = torch.zeros(len(r_ids))
+                    mask[-1] = 1  # 最后一个token有真实奖励
+                    masks_list.append(mask)
+
+            # 转换为tensor
+            old_log_probs = torch.cat(old_log_probs_list).to(self.device)
+            values = torch.cat(values_list).to(self.device)
+            masks = torch.cat(masks_list).to(self.device)
+            rewards = torch.zeros_like(values).to(self.device)
+
+            # 设置奖励（只在最后一个token加上环境奖励）
+            for i, (r, mask) in enumerate(zip(rewards_list, masks_list)):
+                if mask[-1] > 0:
+                    # KL惩罚
+                    kl = old_log_probs[i] - old_log_probs[i]  # 这里简化了，实际要用ref_model
+                    kl_penalty = -self.args.kl_beta * kl
+                    rewards[i] = kl_penalty + r
+
+            # 3. 计算advantages和returns
+            advantages, returns = self.compute_advantages_and_returns(
+                old_log_probs.unsqueeze(0),
+                values.unsqueeze(0),
+                rewards.unsqueeze(0),
+                masks.unsqueeze(0)
+            )
+
+            # 4. PPO多次更新
+            batch_data = (query_ids_list, response_ids_list, old_log_probs,
+                          advantages.squeeze(0), returns.squeeze(0), masks)
+
+            for _ in range(self.args.ppo_epochs):
+                loss, policy_loss, value_loss = self.ppo_step(batch_data)
+
+                self.optimizer.zero_grad()
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(self.actor_critic.parameters(), 1.0)
+                self.optimizer.step()
+
+                total_policy_loss += policy_loss.item()
+                total_value_loss += value_loss.item()
+                num_batches += 1
+                self.training_step += 1
+
+            if batch_idx % 100 == 0:
+                print(f"Batch {batch_idx}/{len(dataset)}: "
+                      f"policy_loss={total_policy_loss / num_batches:.4f}, "
+                      f"value_loss={total_value_loss / num_batches:.4f}")
+
+        return total_policy_loss / num_batches, total_value_loss / num_batches
+
+    def train(self):
+        """主训练循环"""
+        dataset = self.prepare_dataset()
+        print(f"Dataset size: {len(dataset)}")
+
+        for epoch in range(self.args.max_epochs):
+            print(f"\n=== Epoch {epoch + 1}/{self.args.max_epochs} ===")
+            policy_loss, value_loss = self.train_epoch(dataset)
+            print(f"Epoch {epoch + 1} finished: "
+                  f"policy_loss={policy_loss:.4f}, value_loss={value_loss:.4f}")
+
+
+def main():
+    args = get_args()
+    print("PPO Training with CPU")
+    print(f"Arguments: {args}")
+
+    # 创建agent并开始训练
+    agent = PPOAgent(args)
+    agent.train()
+
+    # 保存模型
+    output_dir = Path(args.sft_model_name) / "ppo_trained"
+    output_dir.mkdir(exist_ok=True, parents=True)
+    agent.actor_critic.save_pretrained(output_dir)
+    agent.tokenizer.save_pretrained(output_dir)
+    print(f"Model saved to {output_dir}")
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2/step_5_ppo_rlhf2.py

@@ -0,0 +1,430 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import copy
+import os
+from pathlib import Path
+import platform
+import random
+from typing import Any, Dict, List, Optional, Union, Tuple
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from datasets import load_dataset
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from transformers import (AutoTokenizer,
+                          GPT2PreTrainedModel, GPT2Config, GPT2Model, GPT2LMHeadModel,
+                          )
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from transformers import DataCollatorWithPadding
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--reward_model_name",
+        default=(project_path / "trained_models/gpt2-sst2-reward").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--sft_model_name",
+        default=(project_path / "trained_models/gpt2-sst2-generation").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--dataset_path",
+        default="stanfordnlp/sst2",
+        type=str
+    ),
+    parser.add_argument("--dataset_name", default=None, type=str),
+    parser.add_argument("--dataset_split", default=None, type=str),
+    parser.add_argument(
+        "--dataset_cache_dir",
+        default=(temp_directory / "hub_datasets").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--model_cache_dir",
+        default=(temp_directory / "hub_models").as_posix(),
+        type=str
+    ),
+    parser.add_argument("--dataset_streaming", default=None, type=str),
+    parser.add_argument("--valid_dataset_size", default=1000, type=int),
+    parser.add_argument("--shuffle_buffer_size", default=5000, type=int),
+
+    parser.add_argument(
+        "--output_model_dir",
+        default=(project_path / "trained_models/gpt2-sst2-generation").as_posix(),
+        type=str
+    ),
+    # train
+    parser.add_argument("--batch_size", default=32, type=int)
+
+    # generator
+    parser.add_argument(
+        "--max_new_tokens",
+        default=128, # 8192, 128
+        type=int,
+    )
+    parser.add_argument("--top_p", default=0.85, type=float)
+    parser.add_argument("--temperature", default=0.85, type=float)
+
+    # other
+    parser.add_argument(
+        "--num_workers",
+        default=None if platform.system() in ("Windows", "Darwin") else os.cpu_count() // 2,
+        type=int
+    ),
+    parser.add_argument(
+        "--device",
+        default=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        type=int
+    ),
+    args = parser.parse_args()
+    return args
+
+
+class RewardHead(nn.Module):
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.linear = nn.Linear(self.hidden_size, 1)
+        self._post_init()
+
+    def _post_init(self):
+        nn.init.normal_(
+            self.linear.weight,
+            std=(1.0 / np.sqrt(self.hidden_size + 1))
+        )
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(self, hidden_states):
+        # hidden_states shape: [batch_size, seq_len, hidden_size]
+        reward_logits = self.linear(hidden_states)
+        # reward_logits shape: [batch_size, seq_len, 1]
+        return reward_logits
+
+
+class GPT2RewardModel(GPT2PreTrainedModel):
+    def __init__(self, config: GPT2Config):
+        super().__init__(config)
+        self.transformer = GPT2Model(config)
+        self.reward_head = RewardHead(config.hidden_size)
+        self.post_init()
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> Union[Tuple, torch.Tensor]:
+        transformer_outputs = self.transformer(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True
+        )
+        last_hidden_state = transformer_outputs.hidden_states[-1]
+        # last_hidden_state shape: [batch_size, seq_len, hidden_size]
+        rewards_logits = self.reward_head(last_hidden_state)
+        # rewards_logits shape: [batch_size, seq_len, 1]
+        rewards_logits = torch.squeeze(rewards_logits, -1)
+        # rewards_logits shape: [batch_size, seq_len]
+        rewards = torch.sigmoid(rewards_logits)
+        # rewards shape: [batch_size, seq_len]
+        return rewards
+
+
+class ValueHead(nn.Module):
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.linear = nn.Linear(self.hidden_size, 1)
+        self._post_init()
+
+    def _post_init(self):
+        nn.init.normal_(
+            self.linear.weight,
+            std=(1.0 / np.sqrt(self.hidden_size + 1))
+        )
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(self, hidden_states):
+        # hidden_states shape: [batch_size, seq_len, hidden_size]
+        reward_logits = self.linear(hidden_states)
+        # reward_logits shape: [batch_size, seq_len, 1]
+        return reward_logits
+
+
+class GPT2ActorCriticModel(GPT2PreTrainedModel):
+    def __init__(self, config: GPT2Config):
+        super().__init__(config)
+        self.lm = GPT2LMHeadModel(config)
+        self.value_head = ValueHead(config.hidden_size)
+        self.post_init()
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> Union[Tuple, torch.Tensor]:
+        transformer_outputs = self.lm.forward(
+            input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+        )
+        lm_logits = transformer_outputs.logits
+
+        # values
+        last_hidden_state = transformer_outputs.hidden_states[-1]
+        # last_hidden_state shape: [batch_size, seq_len, hidden_size]
+        values_logits = self.value_head(last_hidden_state)
+        # values_logits shape: [batch_size, seq_len, 1]
+        values = torch.squeeze(values_logits, -1)
+        # values shape: [batch_size, seq_len]
+        values = torch.sigmoid(values)
+        # values shape: [batch_size, seq_len]
+        return lm_logits, values
+
+    @classmethod
+    def from_pretrained(cls, model_name_or_path, *model_args, **kwargs):
+        config = GPT2Config.from_pretrained(model_name_or_path)
+        model = cls(config)
+        pretrained_model = GPT2LMHeadModel.from_pretrained(model_name_or_path)
+        model.lm.load_state_dict(pretrained_model.state_dict(), strict=False)
+        return model
+
+    def generate(self, *args, **kwargs):
+        return self.lm.generate(*args, **kwargs)
+
+
+def masked_mean(values, mask):
+    # 计算带掩码的平均值
+    return (values * mask).sum() / mask.sum()
+
+def masked_var(values, mask):
+    # 计算带掩码的方差
+    mean = masked_mean(values, mask)
+    centred_values = values - mean
+    return masked_mean(centred_values ** 2, mask)
+
+
+def masked_whiten(values, mask):
+    """
+    对数据进行带掩码的白化处理，
+    让有效数据的方差变为1，但均值保持不变
+    """
+    mean, var = masked_mean(values, mask), masked_var(values, mask)
+    whitened = (values - mean) * torch.rsqrt(var + 1e-8)
+    whitened += mean
+    return whitened
+
+
+def main():
+    args = get_args()
+
+    device = torch.device(args.device)
+
+    # reward_model
+    reward_model = GPT2RewardModel.from_pretrained(args.reward_model_name)
+    reward_model = reward_model.to(args.device)
+    reward_model.eval()
+    reward_tokenizer = AutoTokenizer.from_pretrained(args.reward_model_name)
+    reward_tokenizer.pad_token = reward_tokenizer.eos_token
+
+    # actor_critic_model
+    actor_critic_model = GPT2ActorCriticModel.from_pretrained(args.sft_model_name)
+    actor_critic_model = actor_critic_model.to(args.device)
+    actor_critic_tokenizer = AutoTokenizer.from_pretrained(args.sft_model_name)
+    actor_critic_tokenizer.pad_token = actor_critic_tokenizer.eos_token
+    actor_critic_tokenizer.pad_token_id = actor_critic_tokenizer.eos_token_id
+
+    # ref_model
+    ref_model = copy.deepcopy(actor_critic_model)
+    ref_model = ref_model.to(args.device)
+    ref_model.eval()
+
+    dataset_dict = load_dataset(
+        path=args.dataset_path,
+        name=args.dataset_name,
+        split=args.dataset_split,
+        cache_dir=args.dataset_cache_dir,
+        # num_proc=args.num_workers if not args.dataset_streaming else None,
+        streaming=args.dataset_streaming,
+    )
+    train_dataset = dataset_dict["train"]
+    # valid_dataset = dataset_dict["validation"]
+    # test_dataset = dataset_dict["test"]
+
+    def format_func(example):
+        sentence: str = example["sentence"]
+        score: float = float(example["label"])
+        tokenized = actor_critic_tokenizer(sentence)
+        input_ids = tokenized["input_ids"]
+        attention_mask = tokenized["attention_mask"]
+        result = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+        }
+        return result
+
+    train_dataset = train_dataset.map(
+        format_func,
+        batched=False,
+        remove_columns=train_dataset.column_names,
+    )
+    train_dataset = train_dataset.filter(
+        function=lambda x: len(x["input_ids"]) > 8
+    )
+    def token_truncate(example):
+        target_length = random.randint(2, 6)
+        input_ids = example["input_ids"]
+        attention_mask = example["attention_mask"]
+        input_ids = input_ids[:target_length]
+        attention_mask = attention_mask[:target_length]
+        text = actor_critic_tokenizer.decode(input_ids)
+        result = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            # "text": text,
+        }
+        return result
+
+    train_dataset = train_dataset.map(
+        token_truncate,
+        batched=False,
+        remove_columns=train_dataset.column_names,
+    )
+    data_collator = DataCollatorWithPadding(
+        tokenizer=actor_critic_tokenizer,
+        padding=True,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers or 0,
+        collate_fn=data_collator,
+    )
+
+    # for train_batch in train_data_loader:
+    #     print(train_batch)
+
+    for epoch_id in range(10):
+        for batch in train_data_loader:
+            input_ids = batch["input_ids"]
+            attention_mask = batch["attention_mask"]
+
+            query_ids_list = list()
+            query_and_response_ids_list = list()
+            response_ids_list = list()
+            reward_list = list()
+            for idx in range(args.batch_size):
+                input_ids_ = input_ids[idx]
+                attention_mask_ = attention_mask[idx]
+                input_ids_ = input_ids_.to(device)
+                attention_mask_ = attention_mask_.to(device)
+
+                with torch.no_grad():
+                    query_and_response_ids = actor_critic_model.generate(
+                        input_ids=input_ids_.unsqueeze(0),
+                        attention_mask=attention_mask_.unsqueeze(0),
+                        max_new_tokens=random.randint(5, 16),
+                        do_sample=True,
+                        top_p=0.85,
+                        temperature=0.85,
+                        pad_token_id=actor_critic_tokenizer.pad_token_id,
+                        eos_token_id=actor_critic_tokenizer.eos_token_id,
+                        repetition_penalty=1.0,
+                        early_stopping=True,
+                    ).squeeze(0)
+                    query_ids_list.append(input_ids_)
+                    query_and_response_ids_list.append(query_and_response_ids)
+                    response_ids = query_and_response_ids[len(input_ids_):]
+                    response_ids_list.append(response_ids)
+
+                    reward = reward_model(
+                        input_ids=query_and_response_ids.unsqueeze(0),
+                        attention_mask=torch.ones_like(query_and_response_ids, dtype=torch.long).unsqueeze(0),
+                    ).squeeze(0)[-1]
+                    # 将奖励模型的评分从(0,1)缩放到(-1,1)
+                    reward = 2 * (reward - 0.5)
+                    reward_list.append(reward)
+
+            for query_ids, query_and_response_ids in zip(query_ids_list, query_and_response_ids_list):
+                print(actor_critic_tokenizer.decode(query_ids, skip_special_tokens=False))
+                print(actor_critic_tokenizer.decode(query_ids, skip_special_tokens=True))
+                print(actor_critic_tokenizer.decode(query_and_response_ids, skip_special_tokens=True))
+            exit(0)
+            #计算奖励
+            batch_ = list()
+            for query_and_response_ids in query_and_response_ids_list:
+                print(actor_critic_tokenizer.decode(query_and_response_ids))
+                batch_.append({
+                    "input_ids": query_and_response_ids,
+                    "attention_mask": torch.ones_like(query_and_response_ids),
+                })
+
+            batch_ = data_collator(batch_)
+            input_ids = batch_["input_ids"]
+            attention_mask = batch_["attention_mask"]
+            input_ids = input_ids.to(device)
+            attention_mask = attention_mask.to(device)
+
+            logits, values = actor_critic_model(input_ids=input_ids, attention_mask=attention_mask)
+            ref_logits, _ = ref_model(input_ids=input_ids, attention_mask=attention_mask)
+            log_prob = torch.nn.functional.log_softmax(logits[:, :-1, :], dim=-1)
+            ref_log_prob = torch.nn.functional.log_softmax(ref_logits[:, :-1, :], dim=-1)
+            index = input_ids[:, 1:].unsqueeze(-1)
+            log_prob = torch.gather(log_prob, dim=2, index=index).squeeze(-1)
+            ref_log_prob = torch.gather(ref_log_prob, dim=2, index=index).squeeze(-1)
+
+            kl = log_prob - ref_log_prob
+            beta = 0.2
+            kl_penalty = - beta * kl
+
+            rewards = kl_penalty
+            masks = torch.zeros_like(input_ids[:, 1:])
+            for idx in range(args.batch_size):
+                start = len(query_ids_list[idx]) - 1
+                end = start + len(response_ids_list[idx])
+                masks[idx, start:end] = 1
+                rewards[idx, end - 1] += reward_list[idx]
+                values[idx, :-1] *= masks[idx, :]
+                values[idx, -1] = 0
+            rewards = rewards * masks
+
+            # log_prob, rewards, kl_penalty, masks shape: [b, seq_len - 1]
+            # values shape: [b, seq_len]
+
+            # 计算优势
+            seq_len = rewards.shape[-1]
+            last_gae = 0.0
+            gamma, lam = 1.0, 0.95
+            advantage_reversed = list()
+            for t in reversed(range(seq_len)):
+                next_value = values[:, t + 1] if t < seq_len - 1 else 0.0
+                delta = rewards[:, t] + gamma * next_value - values[:, t]
+                last_gae = delta + gamma * lam * last_gae
+                advantage_reversed.append(last_gae)
+            advantages = torch.stack(advantage_reversed[::-1], dim=1)
+            advantages = masked_whiten(advantages, masks)
+            returns = advantages + values[:, :-1]
+
+            # advantages shape: [b, seq_len-1]
+            # returns shape: [b, seq_len-1]
+
+            exit(0)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/rlhf/gpt2_sst2/step_5_pre_ppo_rlhf.py

@@ -0,0 +1,257 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import os
+from pathlib import Path
+import platform
+from typing import Any, Dict, List, Optional, Union, Tuple
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from datasets import load_dataset
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from transformers import (AutoTokenizer,
+                          GPT2PreTrainedModel, GPT2Config, GPT2Model, GPT2LMHeadModel,
+                          )
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--reward_model_name",
+        default=(project_path / "trained_models/gpt2-sst2-reward").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--policy_model_name",
+        default=(project_path / "trained_models/gpt2-sst2-generation").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--dataset_path",
+        default="stanfordnlp/sst2",
+        type=str
+    ),
+    parser.add_argument("--dataset_name", default=None, type=str),
+    parser.add_argument("--dataset_split", default=None, type=str),
+    parser.add_argument(
+        "--dataset_cache_dir",
+        default=(temp_directory / "hub_datasets").as_posix(),
+        type=str
+    ),
+    parser.add_argument(
+        "--model_cache_dir",
+        default=(temp_directory / "hub_models").as_posix(),
+        type=str
+    ),
+    parser.add_argument("--dataset_streaming", default=None, type=str),
+    parser.add_argument("--valid_dataset_size", default=1000, type=int),
+    parser.add_argument("--shuffle_buffer_size", default=5000, type=int),
+
+    parser.add_argument(
+        "--output_model_dir",
+        default=(project_path / "trained_models/gpt2-sst2-generation").as_posix(),
+        type=str
+    ),
+
+    parser.add_argument(
+        "--max_new_tokens",
+        default=128, # 8192, 128
+        type=int, help="最大生成长度（注意：并非模型实际长文本能力）"
+    )
+    parser.add_argument("--top_p", default=0.85, type=float, help="nucleus采样阈值（0-1）")
+    parser.add_argument("--temperature", default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
+
+    # other
+    parser.add_argument(
+        "--num_workers",
+        default=None if platform.system() in ("Windows", "Darwin") else os.cpu_count() // 2,
+        type=int
+    ),
+    parser.add_argument(
+        "--device",
+        default=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        type=int
+    ),
+    args = parser.parse_args()
+    return args
+
+
+class RewardHead(nn.Module):
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.linear = nn.Linear(self.hidden_size, 1)
+        self._post_init()
+
+    def _post_init(self):
+        nn.init.normal_(
+            self.linear.weight,
+            std=(1.0 / np.sqrt(self.hidden_size + 1))
+        )
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(self, hidden_states):
+        # hidden_states shape: [batch_size, seq_len, hidden_size]
+        reward_logits = self.linear(hidden_states)
+        # reward_logits shape: [batch_size, seq_len, 1]
+        return reward_logits
+
+
+class GPT2RewardModel(GPT2PreTrainedModel):
+    def __init__(self, config: GPT2Config):
+        super().__init__(config)
+        self.transformer = GPT2Model(config)
+        self.reward_head = RewardHead(config.hidden_size)
+        self.post_init()
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> Union[Tuple, torch.Tensor]:
+        transformer_outputs = self.transformer(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True
+        )
+        last_hidden_state = transformer_outputs.hidden_states[-1]
+        # last_hidden_state shape: [batch_size, seq_len, hidden_size]
+        rewards_logits = self.reward_head(last_hidden_state)
+        # rewards_logits shape: [batch_size, seq_len, 1]
+        rewards_logits = torch.squeeze(rewards_logits, -1)
+        # rewards_logits shape: [batch_size, seq_len]
+        rewards = torch.sigmoid(rewards_logits)
+        # rewards shape: [batch_size, seq_len]
+        return rewards
+
+
+class ValueHead(nn.Module):
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.linear = nn.Linear(self.hidden_size, 1)
+        self._post_init()
+
+    def _post_init(self):
+        nn.init.normal_(
+            self.linear.weight,
+            std=(1.0 / np.sqrt(self.hidden_size + 1))
+        )
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(self, hidden_states):
+        # hidden_states shape: [batch_size, seq_len, hidden_size]
+        reward_logits = self.linear(hidden_states)
+        # reward_logits shape: [batch_size, seq_len, 1]
+        return reward_logits
+
+
+class GPT2ActorCriticModel(GPT2PreTrainedModel):
+    def __init__(self, config: GPT2Config):
+        super().__init__(config)
+        self.lm = GPT2LMHeadModel(config)
+        self.value_head = ValueHead(config.hidden_size)
+        self.post_init()
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> Union[Tuple, torch.Tensor]:
+        transformer_outputs = self.lm.forward(
+            input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+        )
+        lm_logits = transformer_outputs.logits
+
+        # values
+        last_hidden_state = transformer_outputs.hidden_states[-1]
+        # last_hidden_state shape: [batch_size, seq_len, hidden_size]
+        values_logits = self.value_head(last_hidden_state)
+        # values_logits shape: [batch_size, seq_len, 1]
+        values = torch.squeeze(values_logits, -1)
+        # values shape: [batch_size, seq_len]
+        values = torch.sigmoid(values)
+        # values shape: [batch_size, seq_len]
+        return lm_logits, values
+
+    @classmethod
+    def from_pretrained(cls, model_name_or_path, *model_args, **kwargs):
+        config = GPT2Config.from_pretrained(model_name_or_path)
+        model = cls(config)
+        pretrained_model = GPT2LMHeadModel.from_pretrained(model_name_or_path)
+        model.lm.load_state_dict(pretrained_model.state_dict(), strict=False)
+        return model
+
+    def generate(self, *args, **kwargs):
+        return self.lm.generate(*args, **kwargs)
+
+
+def main():
+    args = get_args()
+
+    reward_model = GPT2RewardModel.from_pretrained(args.reward_model_name)
+    reward_model = reward_model.to(args.device)
+    reward_tokenizer = AutoTokenizer.from_pretrained(args.reward_model_name)
+    reward_tokenizer.pad_token = reward_tokenizer.eos_token
+    print(reward_model)
+    print(reward_tokenizer)
+
+    # tokenized = reward_tokenizer(
+    #     "this is very good movie, I recommend it.",
+    #     return_tensors="pt"
+    # )
+    # rewards = reward_model(**tokenized)
+    # rewards = rewards[0]
+    # rewards = rewards.detach().cpu().numpy()
+    # last_token_reward = rewards[-1]
+    # # rewards: {rewards}\n
+    # msg = f"last_token_reward: {last_token_reward}\n"
+    # print(msg)
+    # exit(0)
+
+    # actor_critic_model
+    actor_critic_model = GPT2ActorCriticModel.from_pretrained(args.policy_model_name)
+    actor_critic_model = actor_critic_model.to(args.device)
+    actor_critic_tokenizer = AutoTokenizer.from_pretrained(args.policy_model_name)
+    actor_critic_tokenizer.pad_token = actor_critic_tokenizer.eos_token
+    print(actor_critic_model)
+    print(actor_critic_tokenizer)
+
+    tokenized = actor_critic_tokenizer(
+        "this is ",
+        return_tensors="pt"
+    )
+    streamer = TextStreamer(actor_critic_tokenizer, skip_prompt=True, skip_special_tokens=True)
+    generated_ids = actor_critic_model.generate(
+        inputs=tokenized["input_ids"], attention_mask=tokenized["attention_mask"],
+        max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
+        pad_token_id=actor_critic_tokenizer.pad_token_id, eos_token_id=actor_critic_tokenizer.eos_token_id,
+        top_p=args.top_p, temperature=args.temperature, repetition_penalty=1.0,
+    )
+    response = actor_critic_tokenizer.decode(generated_ids[0][len(tokenized["input_ids"][0]):], skip_special_tokens=True)
+    print(response)
+
+    tokenized = actor_critic_tokenizer(
+        "this is very good movie, I recommend it.",
+        return_tensors="pt"
+    )
+    lm_logits, values = actor_critic_model(**tokenized)
+    print(values)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

tabs/chat_template_tab.py

@@ -13,8 +13,10 @@ def run_chat_template(conversation: str, model_name: str, add_generation_prompt:
 
     result = tokenizer.apply_chat_template(
         conversation,
+        # tools=None,
         tokenize=False,
         add_generation_prompt=add_generation_prompt,
+        # enable_thinking=True,
     )
     return result