q-transformer

Google Deepmind 提出的 Q-Transformer 实现，一种可扩展的离线强化学习方法，通过自回归 Q 函数实现。

为了最终与提出的多动作自回归 Q 学习进行比较，我将保留单一动作 Q 学习的逻辑。这也可以作为我自己和公众的教育资料。

安装

$ pip install q-transformer

使用方法

import torch

from q_transformer import (
    QRoboticTransformer,
    QLearner,
    Agent,
    ReplayMemoryDataset
)

# 注意力模型

model = QRoboticTransformer(
    vit = dict(
        num_classes = 1000,
        dim_conv_stem = 64,
        dim = 64,
        dim_head = 64,
        depth = (2, 2, 5, 2),
        window_size = 7,
        mbconv_expansion_rate = 4,
        mbconv_shrinkage_rate = 0.25,
        dropout = 0.1
    ),
    num_actions = 8,
    action_bins = 256,
    depth = 1,
    heads = 8,
    dim_head = 64,
    cond_drop_prob = 0.2,
    dueling = True
)

# 你需要通过重写 BaseEnvironment 来提供自己的环境

from q_transformer.mocks import MockEnvironment

env = MockEnvironment(
    state_shape = (3, 6, 224, 224),
    text_embed_shape = (768,)
)

# env.init()     应返回指令和初始状态：Tuple[str, Tensor[*state_shape]]
# env(actions)   应返回奖励、下一个状态和完成标志：Tuple[Tensor[()], Tensor[*state_shape], Tensor[()]]

# agent 是一个类，允许 q-model 与环境交互，生成用于学习的回放记忆数据集

agent = Agent(
    model,
    environment = env,
    num_episodes = 1000,
    max_num_steps_per_episode = 100,
)

agent()

# 在回放记忆数据集上对模型进行 Q 学习

q_learner = QLearner(
    model,
    dataset = ReplayMemoryDataset(),
    num_train_steps = 10000,
    learning_rate = 3e-4,
    batch_size = 4,
    grad_accum_every = 16,
)

q_learner()

# 经过大量学习后
# 你的机器人应该能更好地选择最优动作

video = torch.randn(2, 3, 6, 224, 224)

instructions = [
    '把桌子上的那个苹果拿给我',
    '请把黄油递过来'
]

actions = model.get_optimal_actions(video, instructions)

致谢

感谢 StabilityAI、A16Z 开源 AI 资助计划和 🤗 Huggingface 的慷慨赞助，以及我的其他赞助商，让我能够独立地开源当前的人工智能研究。

待办事项

引用

@inproceedings{qtransformer,
    title   = {Q-Transformer: Scalable Offline Reinforcement Learning via Autoregressive Q-Functions},
    authors = {Yevgen Chebotar and Quan Vuong and Alex Irpan and Karol Hausman and Fei Xia and Yao Lu and Aviral Kumar and Tianhe Yu and Alexander Herzog and Karl Pertsch and Keerthana Gopalakrishnan and Julian Ibarz and Ofir Nachum and Sumedh Sontakke and Grecia Salazar and Huong T Tran and Jodilyn Peralta and Clayton Tan and Deeksha Manjunath and Jaspiar Singht and Brianna Zitkovich and Tomas Jackson and Kanishka Rao and Chelsea Finn and Sergey Levine},
    booktitle = {7th Annual Conference on Robot Learning},
    year   = {2023}
}

@inproceedings{dao2022flashattention,
    title   = {闪光注意力：具有IO感知的快速且内存高效的精确注意力机制},
    author  = {陶, 崔 和 傅, 丹尼尔 Y. 和 尔蒙, 斯特凡诺 和 鲁德拉, 阿特里 和 雷, 克里斯托弗},
    booktitle = {神经信息处理系统进展},
    year    = {2022}
}