使用预训练模型¶

本教程介绍如何在 TorchRL 中使用预训练模型。

import tempfile

完成本教程后，你将能够使用预训练模型进行高效的图像表示，并对其进行微调。

TorchRL 提供了预训练模型，这些模型既可以作为 transforms 使用，也可以作为策略的组成部分。由于它们的语义相同，因此可以在不同上下文中互换使用。在本教程中，我们将使用 R3M (https://arxiv.org/abs/2203.12601)，但其他模型（例如 VIP）也同样适用。

import torch.cuda
from tensordict.nn import TensorDictSequential
from torch import nn
from torchrl.envs import Compose, R3MTransform, TransformedEnv
from torchrl.envs.libs.gym import GymEnv
from torchrl.modules import Actor

is_fork = multiprocessing.get_start_method() == "fork"
device = (
    torch.device(0)
    if torch.cuda.is_available() and not is_fork
    else torch.device("cpu")
)

首先让我们创建一个环境。为简单起见，我们将使用一个常见的 gym 环境。在实践中，这在更具挑战性的具身 AI 场景（例如，查看我们的 Habitat 封装器）中也会有效。

base_env = GymEnv("Ant-v4", from_pixels=True, device=device)

让我们获取预训练模型。我们通过设置 download=True 标志来请求模型的预训练版本。默认情况下此功能是关闭的。接下来，我们将 transform 附加到环境中。实际上，每次收集的数据批次都将通过 transform 并映射到输出 tensordict 中的“r3m_vec”条目。我们的策略，由单层 MLP 组成，将读取此向量并计算相应的动作。

r3m = R3MTransform(
    "resnet50",
    in_keys=["pixels"],
    download=False,  # Turn to true for real-life testing
)
env_transformed = TransformedEnv(base_env, r3m)
net = nn.Sequential(
    nn.LazyLinear(128, device=device),
    nn.Tanh(),
    nn.Linear(128, base_env.action_spec.shape[-1], device=device),
)
policy = Actor(net, in_keys=["r3m_vec"])

让我们检查策略的参数数量

print("number of params:", len(list(policy.parameters())))

我们收集 32 步的 rollout 并打印其输出

rollout = env_transformed.rollout(32, policy)
print("rollout with transform:", rollout)

对于微调，我们在使参数可训练后将 transform 集成到策略中。实际上，更明智的做法可能是将其限制在参数的子集（例如 MLP 的最后一层）。

r3m.train()
policy = TensorDictSequential(r3m, policy)
print("number of params after r3m is integrated:", len(list(policy.parameters())))

再次，我们使用 R3M 收集 rollout。输出的结构略有变化，因为现在环境返回像素（而不是 embedding）。embedding“r3m_vec”是我们策略的中间结果。

rollout = base_env.rollout(32, policy)
print("rollout, fine tuning:", rollout)

我们之所以能如此轻松地将 transform 从环境切换到策略，是因为它们都表现得像 TensorDictModule：它们有一组 “in_keys” 和 “out_keys”，这使得在不同上下文中读取和写入输出变得容易。

作为本教程的总结，让我们看看如何使用 R3M 读取存储在回放缓冲区（例如，在离线 RL 场景中）中的图像。首先，让我们构建数据集

from torchrl.data import LazyMemmapStorage, ReplayBuffer

buffer_scratch_dir = tempfile.TemporaryDirectory().name
storage = LazyMemmapStorage(1000, scratch_dir=buffer_scratch_dir)
rb = ReplayBuffer(storage=storage, transform=Compose(lambda td: td.to(device), r3m))

现在我们可以收集数据（为演示目的使用随机 rollouts）并用它填充回放缓冲区

total = 0
while total < 1000:
    tensordict = base_env.rollout(1000)
    rb.extend(tensordict)
    total += tensordict.numel()

让我们检查回放缓冲区的存储内容。它应该不包含“r3m_vec”条目，因为我们还没有使用它

print("stored data:", storage._storage)

采样时，数据将通过 R3M transform，得到我们想要的处理后的数据。通过这种方式，我们可以在由图像组成的数据集上离线训练算法

batch = rb.sample(32)
print("data after sampling:", batch)

由 Sphinx-Gallery 生成的图库

使用预训练模型¶

文档

教程

资源