前言

过年期间，deepseek-R1火出了圈，各家媒体都在狂轰乱炸宣传deepseek，以及技术圈各种文章来介绍其实现原理。那这里仅从“高质量数据”角度作为入口来阐述对其的理解。

为什么需要高质量数据集？

这不废话吗，没有高质量数据集怎么训练高质量模型。对的，这个回答完全没有问题，从目前大模型能力角度来讲，其回答已经持平或者某些方面已经高于绝大多数人的认知，这是其一。其二是在垂直领域或者具体业务，更多、更贴和的真实数据会给模型带来更好的效果以及降低一个量级的参数量，这也是下游能够应用大模型的主要原因。

那其三呢，我们无法构建一个真实环境，来给大模型进行交互，让其不断试错和学习，所以我们需要将人类对于各种问题的理解以及自然规律等等，用其简单明了直接的方式告诉大模型，这个就是答案,所以出现了SFT。其他文章在介绍到SFT时，喜欢用与人类能够交互的方式来解释这里，我觉得这里的原因是同等重要的，如果模型能够进化到具身智能，那么人类知识作为其主要部分，仍是绕不开与人类交互的。

什么是高质量数据集？

那自然伸展到这里了，或者说高质量数据集应该包含哪些可能，以及什么样的高质量数据集可以更大可能的激发出大模型的能力（aha moment或者是能力涌现）。aha moment是指经过先思考再回答，think step by step，对于推理问题，例如数学、代码方面表现出了优秀的结果。能力涌现是在pretrained model到sft model,指少量标注的QA可以激发出Pretrained model按照人类交互的方式回答用户问题。至此总结了两种高质量数据集的可能表达形式。这也是OpenAI O1、deepseek目前所带来的新的能力。那多模态大模型和agent交互所带来的，我认为也会是接下来的发展方向。

下面展示一个对话demo,经过这种reason dataset来让经过sft model训练后所带来的思考能力。

介绍

Llava是一个多模态大模型，本文以如下代码大致介绍下。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'  # noqa

import requests
from PIL import Image
from transformers import AutoProcessor, LlavaForConditionalGeneration, LlavaConfig


config = LlavaConfig.from_pretrained("llava-hf/llava-1.5-7b-hf")


model = LlavaForConditionalGeneration(config=config)

processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
prompt = "USER: <image>\nWhat's the content of the image? ASSISTANT:"
url = "https://www.ilankelman.org/stopsigns/australia.jpg"
image = Image.open(requests.get(url, stream=True).raw)
inputs = processor(text=prompt, images=image, return_tensors="pt")

generate_ids = model.generate(**inputs, max_new_tokens=15)
processor.batch_decode(generate_ids, skip_special_tokens=True,
                       clean_up_tokenization_spaces=False)[0]

print()

模型结构

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63

LlavaForConditionalGeneration(
  (vision_tower): CLIPVisionModel(
    (vision_model): CLIPVisionTransformer(
      (embeddings): CLIPVisionEmbeddings(
        (patch_embedding): Conv2d(3, 1024, kernel_size=(14, 14), stride=(14, 14), bias=False)
        (position_embedding): Embedding(577, 1024)
      )
      (pre_layrnorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
      (encoder): CLIPEncoder(
        (layers): ModuleList(
          (0-23): 24 x CLIPEncoderLayer(
            (self_attn): CLIPAttention(
              (k_proj): Linear(in_features=1024, out_features=1024, bias=True)
              (v_proj): Linear(in_features=1024, out_features=1024, bias=True)
              (q_proj): Linear(in_features=1024, out_features=1024, bias=True)
              (out_proj): Linear(in_features=1024, out_features=1024, bias=True)
            )
            (layer_norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
            (mlp): CLIPMLP(
              (activation_fn): QuickGELUActivation()
              (fc1): Linear(in_features=1024, out_features=4096, bias=True)
              (fc2): Linear(in_features=4096, out_features=1024, bias=True)
            )
            (layer_norm2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
          )
        )
      )
      (post_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
    )
  )
  (multi_modal_projector): LlavaMultiModalProjector(
    (linear_1): Linear(in_features=1024, out_features=4096, bias=True)
    (act): GELUActivation()
    (linear_2): Linear(in_features=4096, out_features=4096, bias=True)
  )
  (language_model): LlamaForCausalLM(
    (model): LlamaModel(
      (embed_tokens): Embedding(32064, 4096)
      (layers): ModuleList(
        (0-31): 32 x LlamaDecoderLayer(
          (self_attn): LlamaSdpaAttention(
            (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
            (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
            (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
            (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
            (rotary_emb): LlamaRotaryEmbedding()
          )
          (mlp): LlamaMLP(
            (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
            (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
            (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
            (act_fn): SiLU()
          )
          (input_layernorm): LlamaRMSNorm()
          (post_attention_layernorm): LlamaRMSNorm()
        )
      )
      (norm): LlamaRMSNorm()
    )
    (lm_head): Linear(in_features=4096, out_features=32064, bias=False)
  )
)

一共三部分：

1. CLIPVision负责处理图像部分
2. Llama负责文本部分
3. multi_modal_projector负责将图像hidden_size投影到Llama一样维度。

数据处理

图像部分走的是CLIP处理流程，resize到336*336,所以pixel_values shape为(3, 336, 336)，其他没啥特殊。文字部分走的是Llama,这个就很熟悉了。

visual和text对齐

image走ViT,kernel_size为14，所以计算过程和结果如下：

1
2
3

(336-14)/14+1 = 24
24 * 24 = 596
# 如果考虑CLIP CLS的话就是597。

由于CLIP输出是1024,经过multi_modal_projector后维度为(1,576,4096)，这个也是下面image_features的维度。
至此同一个维度4096。

image插入位置

原prompt如下：

1

prompt = "USER: <image>\nWhat's the content of the image? ASSISTANT:"

<image>是插在指定位置的，那么这里图文对齐和之前的对齐就增添了另外一个含义：不仅要能图文对齐，还要考虑图像插入位置。

参考上图中_merge_input_ids_with_image_features函数以及结合上图中的信息，那么不难得出如下结论：

1
2

(final_embedding[:, 5:576+5, :] == image_features[:, :, :]).all()
# > tensor(True)

由于后续target task为VQA、Image Caption之类的，先到此为止。

前言

最近有个需求，能够对不同排版格式的繁体信息进行抽取，所以从传统的版面分析+文字检测、识别+阅读顺序+NLP到现在发展的多模态大模型综合调研。

此处以文字区域检测+识别做个demo,来直观感受多模态大模型的结果。

总结

1. LVLM(large vision language model)相比LLM至少落后一代。
2. LVLM相比LLM更具有挑战性。1）多模态信息融合。2）从结果上看，训练时长与loss下降速度。3）高清图片，针对不同尺寸的图片，原来例如CLIP使用固长224*224像素来patch,Qwen-VL-Chat使用448,MiniCPM-V-2_6采用动态计算切分方式，来更好贴近原始SigLIP的输入尺寸,减少缩放后图片质量的损失。以及引入query embed来减少高清图片输入长度过长问题，看Perceiver Sampler和LlaVA-UHD。
3. 存在更为明显重复生成问题。之前在试字节豆包时，在一个问答中涉及批量工作场景时遇到过重复生成，即模型“宕机”了，不断重复一段文本。其他LLM也有这类问题，只是相比LVLM更少些。在未见过的或者说更为复杂排版场景时，出错的概率更高。
4. 推理速度。不过这个也是跟随业界大环境走吧，第一代的LLM其实速度也不是很快，后面VLLM等加速推理框架、其他技术出现等。

附件

注：实际比下面还复杂些，此处只做简单demo展示。

简体横排（没有这样训练语料仍能泛化，模型本身具备OCR能力）：

手写体（基本仍符合横排）：

繁体竖排：

对于更复杂场景的排版，泛化效果不理想，官方解释是原始训练语料包含少。

1. graphviz docker

1
2
3
4
5
6
7
8
9
10

FROM ubuntu:18.04

ENV LANG C.UTF-8
WORKDIR /code
COPY . .

RUN sed -i s@/deb.debian.org/@/mirrors.aliyun.com/@g /etc/apt/sources.list && \
 sed -i s@/security.debian.org/@/mirrors.aliyun.com/@g /etc/apt/sources.list && \
 apt-get clean && apt-get update && apt-get install -y python3-pip python3-dev build-essential git graphviz fonts-arphic-ukai

2. networkx

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

from pprint import pprint

import networkx as nx
import matplotlib.pyplot as plt


def list_fonts():
    from matplotlib.font_manager import FontManager
    fm = FontManager()
    pprint(set([f.name for f in fm.ttflist]))


def draw():
    """
    列出系统的字体，然后找一个支持中文的
    """
    plt.rcParams['font.sans-serif'] = ['AR PL UKai CN']
    # 创建一个NetworkX图
    G = nx.Graph()
    G.add_edges_from([('你呀你呀', 2), (2, 3), (3, 1)])

    # 绘制图形
    pos = nx.spring_layout(G)  # 选择布局算法
    nx.draw(G, pos, with_labels=True)  # , node_color='lightblue', edge_color='gray')  # 绘制图形
    plt.savefig('graph.png')  # 保存为PNG图像文件
    plt.show()  # 显示图形（可选）


if __name__ == '__main__':
    # list_fonts()
    draw()

最近面试了一北京候选者，之前使用电子病例以及CT图像两种模态信息，训练ViLT多模态预测模型，提高肺结节良恶性预测准确率。
正好我对多模态如何对齐也比较感兴趣，以Transformers-Tutorials提供的代码为例，来看下其内部是如何实现的。
数据集我没有从VQA下载，太大了，这里也强烈安利huggingface提供的lmms-lab/VQAv2 dataset。

剩下就改下VQADataset部分，其他保持不变。

这里记录比较有趣的几个点。

1. text和image如何对齐？

答案：在第二维对齐。

具体来说，text部分使用的是BertTokenizer,max_position为40（所以如果有长文本，这里就坑了），假设batch_size为4,text embedding出来后就是(4, 40, 768)。
image部分使用的也是ViT,所以整体思想和CLIP是一样的，得到patch后flatten，然后concat到一起即可， 即embeddings = torch.cat([text_embeds, image_embeds], dim=1)。

相比CLIP,不同之处在于支持不同宽高的image，CLIP默认只支持(224, 224)宽高的图像。所以ViLT多了个visual pad操作，另外kernel size也不一样。如果深究的话，可以看看visual_embed部分。

2. 竟然有ViltForTokenClassification？

是的，上面既然已经对齐了，那么前面都是text部分，所以做文本相关例如token classification等这部分也是没问题的。

3. ViltForQuestionAnswering是QA吗？

先入为主来讲，这个A呢，是不固定的，比如使用generation来生成。但是这里呢，是有限的！！！。

怎么说呢，人家提供了个类似[CLS]，然后你可以指定自己的类别，所以他是个分类器。

这里作者使用了finetuned的label来作为演示，可以看看”dandelin/vilt-b32-finetuned-vqa”下config.json中的label2id查看标签数量。

1
2
3
4

from transformers import ViltConfig

config = ViltConfig.from_pretrained("dandelin/vilt-b32-finetuned-vqa")

4. 排序新操作

正如作者所说的那样，VQAv2是个多标签分类数据集，什么意思呢，例如下图。

它的question：where is he looking?

它的标注结果annotations：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

{'question_type': 'none of the above',
 'multiple_choice_answer': 'down',
 'answers': [{'answer': 'down', 'answer_confidence': 'yes', 'answer_id': 1},
  {'answer': 'down', 'answer_confidence': 'yes', 'answer_id': 2},
  {'answer': 'at table', 'answer_confidence': 'yes', 'answer_id': 3},
  {'answer': 'skateboard', 'answer_confidence': 'yes', 'answer_id': 4},
  {'answer': 'down', 'answer_confidence': 'yes', 'answer_id': 5},
  {'answer': 'table', 'answer_confidence': 'yes', 'answer_id': 6},
  {'answer': 'down', 'answer_confidence': 'yes', 'answer_id': 7},
  {'answer': 'down', 'answer_confidence': 'yes', 'answer_id': 8},
  {'answer': 'down', 'answer_confidence': 'yes', 'answer_id': 9},
  {'answer': 'down', 'answer_confidence': 'yes', 'answer_id': 10}],
 'image_id': 262148,
 'answer_type': 'other',
 'question_id': 262148000}

answer是有多个的，也就是说它的answers是个众包的，有的人认为是down,有的人认为是table。

4.1 作者实现思路

1、处理answer

例如上面标注结果，首先统计answer次数，得到下面结果：

1

answers_count = {'down': 7, 'at table': 1, 'skateboard': 1, 'table': 1}

接着对每个answer count使用get_score进行平滑处理，get_score有三种可能，那每种可能就认为是每种答案的权重。

1
2
3
4

def get_score(count: int) -> float:
    return min(1.0, count / 3)

# [0.3333333333333333, 0.6666666666666666, 1.0]

2、ViltForQuestionAnswering

在ViLT出来后，接了个分类器，这个分类器的标签数量来自自己实际业务指定，比如answers的总数量。假设一共有10个类别，down类别是1,at table是2,依次类推，得到label如下：

1

label = [1.0, 0.3333333333333333, 0.3333333333333333, 0.3333333333333333, 0, 0, 0, 0, 0, 0]

接着计算loss如下：

1
2

loss = nn.functional.binary_cross_entropy_with_logits(logits, labels) * labels.shape[1]

这里又乘了个labels总数量，可以简单理解成loss scale。

3、解码

作者使用topk进行解码：

1
2

probs, classes = torch.topk(predicted_classes, 5)

至此，大致理解了整体流程。

这里是get到一种新的实现方式了哈，那接下来我想对其扩展下，看看还有哪些方式可以做这种。

4.2 其他排序思路

1、rank loss。排序中常用到。
2、目前大模型中的打分模型(score model)，在最开始我没看其实现的时候，我曾深深纠结于这又是什么新的技术，等我看完后，我感觉更应该关注的是“打分”俩字,即强化学习思想中的critic model。本质技术层面没有太大变化，作为打分模型为Actor提供优化方向。
3、DPO。对，你没看错，强化学习算法DPO,本质思想就在于好的要逐渐和坏的拉开差距，使其更符合人类偏好。

看的越多，会发现还是有共通之处的。

介绍

分享个模型可解释的库SHAP，其可以对XGBoost等机器学习库进行结果解释。

更多还有类似Paddle中的TrustAI。

什么是SHAP？

SHAP（SHapley Additive exPlanations）是用来解释机器学习模型的工具。它告诉我们每个特征对模型预测结果的贡献有多大。想象一下你和几个朋友一起干活，最后得到了一笔报酬。你想知道每个人应该分多少钱，SHAP就是在做类似的事情，只不过它是在告诉你每个特征对模型预测结果的贡献。

核心概念

1. Shapley值：

想象一个团队在完成任务后得到一笔奖金，我们想公平地分配这笔奖金。Shapley值就是一种分配方法，告诉我们每个成员（在机器学习中是每个特征）应得的报酬（对预测结果的贡献）。

2. 边际贡献：

假设我们在玩一个游戏，每次加入一个新成员（特征），看看他们对团队成绩（模型预测）的贡献有多大。边际贡献就是每次新加入的成员对总成绩的额外贡献。

SHAP如何工作？

SHAP计算每个特征的贡献时，会考虑所有可能的加入顺序。这是因为不同的加入顺序可能会导致不同的贡献。举个例子：

• 如果你和你的朋友们一起完成一个项目，每个朋友加入的顺序不同，他们的贡献也可能不同。
• SHAP通过计算所有可能顺序中的平均贡献来确定每个特征的Shapley值。

具体步骤

1. 考虑所有可能的特征子集：

对于每个特征，计算它加入不同特征子集时，对模型预测结果的边际贡献。

2. 计算平均贡献：

对于每个特征，计算它在所有可能子集中的平均贡献，这就是这个特征的Shapley值。

例子

假设我们有一个模型在预测一个人是否会购买某产品，模型用到了年龄、收入和是否有孩子三个特征。SHAP会做以下事情：

1. 计算每个特征的贡献：

• 先只考虑收入，再加上年龄，看看预测结果有多大变化。再加上是否有孩子，看看预测结果再变化多少。
• 再换个顺序，先考虑年龄，再加上收入，然后加上是否有孩子。
• 如此继续，考虑所有可能的加入顺序。

2. 计算平均值：

对于每个特征，计算它在所有顺序中的平均贡献，这个平均值就是该特征的Shapley值。

总结

SHAP通过考虑每个特征在不同组合中的贡献，并取平均值，来告诉我们每个特征对模型预测的影响有多大。这样，我们就能清楚地知道模型是如何做出预测的，每个特征在其中起了多大作用。

演示代码：https://github.com/geasyheart/algo/blob/m2/python/shap_demo.ipynb

简介

DPO全称Direct Preference Optimization，它是RLHF算法的一种，相比PPO算法来讲，它只需要actor和ref model,少了critic和reward model。其核心期望为good loss - bad loss越来越大， 这点和排序模型中的rank loss很相似，但是又不希望和ref model偏差太多。

最小实现代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

from copy import deepcopy

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, LlamaForCausalLM, LlamaConfig

torch.manual_seed(0)
# 超参数
beta = 0.1
# 加载模型


# data
prompt_ids = [1, 2, 3, 4, 5, 6]
good_response_ids = [7, 8, 9, 10]
# 对loss稍加修改可以应对一个good和多个bad的情况
bad_response_ids_list = [[1, 2, 3, 0], [4, 5, 6, 0]]

# 转换成模型输入
input_ids = torch.LongTensor(
    [prompt_ids + good_response_ids, *[prompt_ids + bad_response_ids for bad_response_ids in bad_response_ids_list]]
)
# labels 提前做个shift
labels = torch.LongTensor(
    [
        [-100] * len(prompt_ids) + good_response_ids,
        *[[-100] * len(prompt_ids) + bad_response_ids for bad_response_ids in bad_response_ids_list]
    ]
)[:, 1:]
loss_mask = (labels != -100)
labels[labels == -100] = 0


policy_model = LlamaForCausalLM(config=LlamaConfig(vocab_size=1000, num_hidden_layers=1, hidden_size=128))
reference_model = deepcopy(policy_model)


# 计算 policy model的log prob
logits = policy_model(input_ids)["logits"][:, :-1, :]
per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)
all_logps = (per_token_logps * loss_mask).sum(-1)
# 暂时写死第一个是good response的概率
policy_good_logps, policy_bad_logps = all_logps[:1], all_logps[1:]

# 计算 reference model的log prob
with torch.no_grad():
    logits = reference_model(input_ids)["logits"][:, :-1, :]
    per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)
    all_logps = (per_token_logps * loss_mask).sum(-1)
    # 暂时写死第一个是good response的概率
    reference_good_logps, reference_bad_logps = all_logps[:1], all_logps[1:]

# 计算loss，会自动进行广播
logits = (policy_good_logps - reference_good_logps) - (policy_bad_logps - reference_bad_logps)
loss = -F.logsigmoid(beta * logits).mean()
print(loss)

参考资料

1. 知乎介绍1
2. trl dpo loss计算

简介

值函数希望学习一个价值函数，这个值可以用于评估当前决策的分值。策略函数希望学习一个策略函数，拿到其动作的概率分布。

Actor-Critic是在策略函数的基础上，额外引入学习价值函数，来帮助策略函数更好地学习。

下面这个图很好表示了两者关系。

重点看actor-critic算法中update函数log_probs部分。actor采用策略，critic来进行评价。

前言

这个是RLHF系列中的策略梯度部分，在看了Hands-on-RL和parl两者实现后，感觉整体难度并不是很高，但是当自己从零实现时还是会莫名其妙多一些问题，相比深度学习来讲，还是有蛮多小细节是需要额外注意的。

注意点

1. log平滑

这里是指learn阶段中的获取最大期望阶段，如下代码所示：

1
2

output = self.model(obs_bs)
output = torch.log(output.gather(-1, action_bs.reshape(-1, 1)))

在最开始自己实现时，我没有加log进行平滑，发现模型没法收敛(CarPole-v0 reward最大得分为200)，一直是8,9徘徊。后来我看了上述实现，发现这里多了个log,这里让我觉得很困惑，因为我觉得这一步是不必要的，原因有以下几个方面：

1. model.forward部分，已经用softmax做归一化了，已经避免了差异较大的情况。
2. 从某种角度来讲，我甚至觉得model.forward中的softmax部分也不应该添加。因为本质来讲就是希望期望最大嘛。

但是呢，如果不加log这一步，模型就无法收敛。

2. 折扣因子

这里是指在每一次done，产生了一批state、reward、action之后，在进行计算loss时，下一步的reward还要考虑当前步reward的结果，即下一步的reward一定要小于当前步的reward。也就是calc_reward_to_go函数这里。

这里同样也会觉得很困惑，因为如果希望期望最大，那就reward * prob使其概率最大即可。

如果看Hands-on-RL他的实现，不会感到任何困惑，因为他是用for来做的。但是呢，parl的实现在考虑了reward * prob使其概率最大这一步之后，又添加了calc_reward_to_go函数，从而本来是个离散的东西，强生生的给变成了一个连续状态的事情，关键呢，怎么看都不像是连续的，因为当前步在计算loss时也并没有跟上一步扯上直接的关系。

这一步还好，从理解角度来讲，我会更倾向Hands-on-RL的实现，容易理解。

总结

关于第一点，我的感觉是步子不能迈太大，宁可慢慢收敛，也要比无法收敛更强，例如model部分我尝试改成如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

class PolicyNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNet, self).__init__()
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        # return F.softmax(self.fc2(x), dim=-1)
        return self.fc2(x)
    @torch.no_grad()
    def sample(self, obs):
        result = self(torch.FloatTensor(obs))
        result = F.softmax(result, dim=-1)
        return torch.multinomial(result, 1).item()

即forward部分不用softmax,剩下代码保持不变，也同样处于无法收敛状态，似乎来看,softmax+log才是这个算法成功的关键。

这里让我想到一个事情，在深度学习Layer参数初始化的过程，比如：

1

a = nn.Linear(2222, 11)

我们会是这么写，基本不会关心a这个linear的参数是如何初始化的。是因为框架内部已经考虑了kaiming、xavier、uniform等各种初始化技术。如果不加这些参数初始化技术，模型基本也很难收敛。

所以如果应用的话，可以采用现成的实现，如果研究的话，其中一些细节可以慢慢调整。

源码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

import random

import gym
import torch
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
import rl_utils
from parl.env import CompatWrapper


class PolicyNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNet, self).__init__()
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)

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

    @torch.no_grad()
    def sample(self, obs):
        result = self(torch.FloatTensor(obs))
        return torch.multinomial(result, 1).item()

def calc_reward_to_go(reward_list, gamma=1.0):
    for i in range(len(reward_list) - 2, -1, -1):
        # G_i = r_i + γ·G_i+1
        reward_list[i] += gamma * reward_list[i + 1]  # Gt
    return np.array(reward_list)

class Agent(object):
    def __init__(self, state_dim, action_dim, hidden_dim):
        self.model = PolicyNet(state_dim=state_dim, hidden_dim=hidden_dim, action_dim=action_dim)
        self.optim = torch.optim.Adam(self.model.parameters(), lr=1e-3)

    def train(self, env):
        total_reward = 0
        obs_bs, action_bs, next_obs_bs, reward_bs, done_bs = [], [], [], [], []
        obs = env.reset()
        while True:
            action = self.model.sample(obs=obs)
            next_obs, reward, done, info = env.step(action=action)
            if done:
                break
            obs_bs.append(obs)
            action_bs.append(action)
            next_obs_bs.append(next_obs)
            reward_bs.append(reward)
            done_bs.append(done)

            total_reward += reward
            obs = next_obs
        if not obs_bs: return total_reward
        # learn
        reward_bs = calc_reward_to_go(reward_bs)

        obs_bs = torch.FloatTensor(np.array(obs_bs))
        action_bs = torch.LongTensor(np.array(action_bs))
        reward_bs = torch.FloatTensor(np.array(reward_bs))

        output = self.model(obs_bs)
        output = torch.log(output.gather(-1, action_bs.reshape(-1, 1)))

        loss = torch.mean(-1 * output.flatten() * reward_bs)

        self.optim.zero_grad()
        loss.backward()
        self.optim.step()

        return total_reward


if __name__ == '__main__':
    env = CompatWrapper(gym.make('CartPole-v0'))
    agent = Agent(env.observation_space.shape[0], action_dim=env.action_space.n, hidden_dim=64)
    for epoch in range(10000):
        total_reward = agent.train(env=env)
        print(f'Epoch[{epoch}]-->{total_reward}')