RayCollector¶
- class torchrl.collectors.distributed.RayCollector(create_env_fn: Union[Callable, EnvBase, List[Callable], List[EnvBase]], policy: Callable[[TensorDict], TensorDict], *, frames_per_batch: int, total_frames: int = -1, device: torch.device | List[torch.device] = None, storing_device: torch.device | List[torch.device] = None, env_device: torch.device | List[torch.device] = None, policy_device: torch.device | List[torch.device] = None, max_frames_per_traj=-1, init_random_frames=-1, reset_at_each_iter=False, postproc=None, split_trajs=False, exploration_type=InteractionType.RANDOM, collector_class: Callable[[TensorDict], TensorDict] = <class 'torchrl.collectors.collectors.SyncDataCollector'>, collector_kwargs: Union[Dict, List[Dict]] = None, num_workers_per_collector: int = 1, sync: bool = False, ray_init_config: Dict = None, remote_configs: Union[Dict, List[Dict]] = None, num_collectors: int = None, update_after_each_batch=False, max_weight_update_interval=-1)[源代码]¶
使用 Ray 后端的分布式数据收集器。
此 Python 类作为基于 Ray 的解决方案，用于在分布式集群中实例化和协调多个数据收集器。与 TorchRL 非分布式收集器类似，此收集器是一个可迭代对象，它会生成 TensorDict，直到达到收集到的帧的目标数量，但在后台处理分布式数据收集。
类字典输入参数“ray_init_config”可用于提供调用 Ray 初始化方法 ray.init() 的关键字参数。如果未提供“ray_init_config”，则默认行为是自动检测现有的 Ray 集群，或者如果未找到现有的集群，则在本地启动一个新的 Ray 实例。有关高级初始化关键字参数，请参阅 Ray 文档。
类似地，字典输入参数“remote_configs”可用于指定调用 ray.remote() 时创建每个远程收集器 actor 的关键字参数，包括收集器计算资源。所有收集器资源的总和应在集群中可用。有关 ray.remote() 方法的高级配置，请参阅 Ray 文档。默认关键字参数为
>>> kwargs = { ... "num_cpus": 1, ... "num_gpus": 0.2, ... "memory": 2 * 1024 ** 3, ... }
收集器实例之间的协调可以指定为“同步”或“异步”。在同步协调中，此类会等待所有远程收集器收集一个 rollout，将所有 rollout 连接到单个 TensorDict 实例中，最后生成连接后的数据。另一方面，如果要异步执行协调，则此类会根据单个远程收集器提供的 rollout。
- 参数：
create_env_fn (Callable 或 List[Callabled]) – 可调用的列表，每个可调用返回
EnvBase的实例。policy (Callable) –
要在环境中执行的策略。必须接受
tensordict.tensordict.TensorDictBase对象作为输入。如果提供None，则使用的策略将是具有环境action_spec的RandomPolicy实例。接受的策略通常是TensorDictModuleBase的子类。这是推荐的收集器用法。也接受其他可调用对象：如果策略不是TensorDictModuleBase（例如，常规的Module实例），它将首先包装在 nn.Module 中。然后，收集器将尝试评估这些模块是否需要包装在TensorDictModule中。- 如果策略前向签名与任何forward(self, tensordict)匹配，forward(self, td)或forward(self, <anything>: TensorDictBase)（或任何以TensorDictBase的子类作为单个参数类型的类型）匹配，则策略不会包装在TensorDictModule中。在所有其他情况下，将尝试对其进行包装，如下所示：
TensorDictModule(policy, in_keys=env_obs_key, out_keys=env.action_keys)。
- 关键字参数：
frames_per_batch (int) – 表示批次中元素总数的仅关键字参数。
total_frames (int, Optional) – 收集器返回的帧总数的下限。一旦帧总数等于或超过传递给收集器的帧总数，迭代器将停止。默认值为 -1，这意味着没有目标帧总数（即收集器将无限期运行）。
device (int, str 或 torch.device, 可选) – 收集器的通用设备。
device参数填充任何未指定的设备：如果device不为None且storing_device、policy_device或env_device中的任何一个未指定，则其值将设置为device。默认为None（无默认设备）。支持设备列表。storing_device (int, str 或 torch.device, 可选) – 输出
TensorDict将存储其上的远程设备。如果传递了device且storing_device为None，则它将默认为device指示的值。对于长轨迹，可能需要将数据存储在与策略和环境执行所在的设备不同的设备上。默认为None（输出 tensordict 不在特定设备上，叶子张量位于创建它们的设备上）。支持设备列表。env_device (int, str 或 torch.device, 可选) – 环境应该被转换到的（或如果支持该功能则执行）远程设备。如果未指定且环境具有非
None设备，则env_device将默认为该值。如果传递了device且env_device=None，则它将默认为device。如果env_device的如此指定的值与policy_device不同，并且其中一个不为None，则数据将在传递给环境之前转换为env_device（即，支持将不同的设备传递给策略和环境）。默认为None。支持设备列表。policy_device (int, str 或 torch.device, 可选) – 策略应该被转换到的远程设备。如果传递了
device且policy_device=None，则它将默认为device。如果policy_device的如此指定的值与env_device不同，并且其中一个不为None，则数据将在传递给策略之前转换为policy_device（即，支持将不同的设备传递给策略和环境）。默认为None。支持设备列表。create_env_kwargs (dict, 可选) –
create_env_fn的关键字参数字典。max_frames_per_traj (int, 可选) – 每条轨迹的最大步数。请注意，一条轨迹可以跨越多个批次（除非
reset_at_each_iter设置为True，请参见下文）。一旦轨迹达到n_steps，环境将被重置。如果环境将多个环境包装在一起，则每个环境的步数将独立跟踪。允许使用负值，在这种情况下，将忽略此参数。默认为None（即，没有最大步数）。init_random_frames (int, 可选) – 策略在被调用之前忽略的帧数。此功能主要用于离线/基于模型的设置，其中可以使用一批随机轨迹来初始化训练。如果提供，它将四舍五入到最接近的 frames_per_batch 的倍数。默认为
None（即，没有随机帧）。reset_at_each_iter (bool, 可选) – 是否应该在批次收集开始时重置环境。默认为
False。postproc (Callable, 可选) – 后处理转换，例如
Transform或MultiStep实例。默认为None。split_trajs (bool, 可选) – 布尔值，指示结果 TensorDict 是否应该根据轨迹进行分割。有关更多信息，请参阅
split_trajectories()。默认为False。exploration_type (ExplorationType, 可选) – 收集数据时要使用的交互模式。必须是
torchrl.envs.utils.ExplorationType.DETERMINISTIC、torchrl.envs.utils.ExplorationType.RANDOM、torchrl.envs.utils.ExplorationType.MODE或torchrl.envs.utils.ExplorationType.MEAN之一。collector_class (Python 类) – 要远程实例化的收集器类。可以是
SyncDataCollector、MultiSyncDataCollector、MultiaSyncDataCollector或这些类的派生类。默认为SyncDataCollector。collector_kwargs (dict 或 list, 可选) – 要传递给远程数据收集器的参数字典。如果提供了列表，则每个元素将对应于专用收集器的一组单独的关键字参数。
num_workers_per_collector (int) – 要在远程节点上使用的 env 构造函数的副本数量。默认为 1（每个收集器一个环境）。在一个工作节点上，所有子工作节点都将执行相同的环境。如果需要执行不同的环境，则应跨工作节点分发它们，而不是子节点。
ray_init_config (dict, 可选) – 用于调用 ray.init() 的关键字参数。
remote_configs (字典列表, 可选) – 每个远程收集器的 ray 资源规范。也可以提供单个字典，它将用于所有收集器。
num_collectors (int, 可选) – 要实例化的收集器总数。
sync (bool) – 如果为
True，则结果 tensordict 是每个节点上收集的所有 tensordict 的堆栈。如果为False（默认值），则每个 tensordict 来自“先到先得”方式的单独节点。update_after_each_batch (bool, 可选) – 如果为
True，则将在每次收集后更新权重。对于sync=True，这意味着所有工作节点都将看到其权重更新。对于sync=False，只有收集数据的那个工作节点的权重会被更新。默认为False，即更新必须通过torchrl.collectors.distributed.RayDistributedCollector.update_policy_weights_()手动执行。max_weight_update_interval (int, 可选) – 在更新工作节点的策略权重之前可以收集的最大批次数。对于同步收集，此参数将被
update_after_each_batch覆盖。对于异步收集，即使打开了update_after_each_batch，也可能某个工作节点在一段时间内没有看到其参数被更新。默认为 -1（无强制更新）。
示例
>>> from torch import nn >>> from tensordict.nn import TensorDictModule >>> from torchrl.envs.libs.gym import GymEnv >>> from torchrl.collectors.collectors import SyncDataCollector >>> from torchrl.collectors.distributed import RayCollector >>> env_maker = lambda: GymEnv("Pendulum-v1", device="cpu") >>> policy = TensorDictModule(nn.Linear(3, 1), in_keys=["observation"], out_keys=["action"]) >>> distributed_collector = RayCollector( ... create_env_fn=[env_maker], ... policy=policy, ... collector_class=SyncDataCollector, ... max_frames_per_traj=50, ... init_random_frames=-1, ... reset_at_each_iter=-False, ... collector_kwargs={ ... "device": "cpu", ... "storing_device": "cpu", ... }, ... num_collectors=1, ... total_frames=10000, ... frames_per_batch=200, ... ) >>> for i, data in enumerate(collector): ... if i == 2: ... print(data) ... break
