RPCDataCollector¶
- class torchrl.collectors.distributed.RPCDataCollector(create_env_fn, policy, *, 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: int = -1, init_random_frames: int = -1, reset_at_each_iter: bool = False, postproc: Callable | None = None, split_trajs: bool = False, exploration_type: ExporationType = InteractionType.RANDOM, collector_class=<class 'torchrl.collectors.collectors.SyncDataCollector'>, collector_kwargs=None, num_workers_per_collector=1, sync=False, slurm_kwargs=None, update_after_each_batch=False, max_weight_update_interval=-1, launcher='submitit', tcp_port=None, visible_devices=None, tensorpipe_options=None)[source]¶
一個基於 RPC 的分散式資料收集器。
支援同步和非同步資料收集。
- 參數:
create_env_fn (Callable 或 List[Callabled]) – Callables 的列表，每個 Callable 回傳一個
EnvBase的實例。policy (Callable) –
要在環境中執行的策略。必須接受
tensordict.tensordict.TensorDictBase物件作為輸入。如果提供None，則使用的策略將是具有環境action_spec的RandomPolicy實例。接受的策略通常是TensorDictModuleBase的子類別。這是建議的收集器使用方式。也接受其他 callables：如果該策略不是TensorDictModuleBase（例如，常規的Module實例），它將首先被包裝在 nn.Module 中。然後，收集器將嘗試評估這些模組是否需要包裝在TensorDictModule中。 - 如果策略 forward 簽章符合forward(self, tensordict),forward(self, td)或forward(self, <anything>: TensorDictBase)（或任何具有單一引數且類型為TensorDictBase子類別的類型註解）中的任何一個，則該策略將不會被包裝在TensorDictModule中。在所有其他情況下，將嘗試將其包裝如下：
TensorDictModule(policy, in_keys=env_obs_key, out_keys=env.action_keys)。
- Keyword Arguments:
frames_per_batch (int) – 一個僅限關鍵字引數，表示批次中的元素總數。
total_frames (int) –
一個僅限關鍵字引數，表示收集器在其生命週期內回傳的總幀數。如果
total_frames無法被frames_per_batch整除，則會引發例外。可以透過傳遞
total_frames=-1來建立無窮的收集器。預設值為-1（無窮的收集器）。device (int, str 或 torch.device, optional) – 收集器的一般裝置。
device引數會填寫任何未指定的裝置：如果device不是None並且storing_device、policy_device或env_device中任何一個未指定，則其值將設定為device。預設值為None（無預設裝置）。支援裝置列表。storing_device (int, str 或 torch.device, optional) – 遠端裝置，輸出
TensorDict將儲存在該裝置上。如果傳遞了device且storing_device為None，則它將預設為device指示的值。對於長軌跡，可能需要將資料儲存在與執行策略和環境的裝置不同的裝置上。預設值為None（輸出 tensordict 不在特定裝置上，葉張量位於它們被建立的裝置上）。支援裝置列表。env_device (int, str 或 torch.device, optional) – 環境應該轉換 (或執行，如果支援此功能) 的遠端裝置。 如果未指定，且環境具有非
None的裝置，則env_device將預設為該值。 如果傳遞了device且env_device=None，則它將預設為device。 如果如此指定的env_device值與policy_device不同，並且它們之中有一個不是None，則資料將在傳遞到環境之前轉換為env_device(即，支援將不同的裝置傳遞給策略和環境)。 預設值為None。 支援裝置列表。policy_device (int, str 或 torch.device, optional) – 策略應該轉換到的遠端裝置。 如果傳遞了
device且policy_device=None，則它將預設為device。 如果如此指定的policy_device值與env_device不同，並且它們之中有一個不是None，則資料將在傳遞到策略之前轉換為policy_device(即，支援將不同的裝置傳遞給策略和環境)。 預設值為None。 支援裝置列表。max_frames_per_traj (int, optional) – 每個軌跡的最大步數。 請注意，一個軌跡可以跨越多個批次（除非
reset_at_each_iter設定為True，請參閱下文）。 一旦軌跡達到n_steps，環境就會重置。 如果環境將多個環境包裝在一起，則會獨立追蹤每個環境的步數。 允許負值，在這種情況下，此參數將被忽略。 預設值為None（即，沒有最大步數）。init_random_frames (int, optional) – 在呼叫策略之前，策略被忽略的幀數。 此功能主要用於離線/基於模型的設定，其中可以使用一批隨機軌跡來初始化訓練。 如果提供，它將四捨五入到 frames_per_batch 的最接近倍數。 預設值為
None（即，沒有隨機幀）。reset_at_each_iter (bool, optional) – 是否應在批次收集開始時重置環境。 預設值為
False。postproc (Callable, optional) – 一個後處理轉換，例如
Transform或MultiStep實例。 預設值為None。split_trajs (bool, optional) – 指示是否應根據軌跡分割結果 TensorDict 的布林值。 有關更多信息，請參閱
split_trajectories()。 預設值為False。exploration_type (ExplorationType, optional) – 收集資料時要使用的互動模式。 必須是
torchrl.envs.utils.ExplorationType.DETERMINISTIC、torchrl.envs.utils.ExplorationType.RANDOM、torchrl.envs.utils.ExplorationType.MODE或torchrl.envs.utils.ExplorationType.MEAN之一。 預設值為torchrl.envs.utils.ExplorationType.RANDOM。collector_class (type 或 str, optional) –
遠端節點的收集器類別。 可以是
SyncDataCollector、MultiSyncDataCollector、MultiaSyncDataCollector或是這些的衍生類別。 字串 "single"、"sync" 和 "async" 對應於各自的類別。 預設值為SyncDataCollector。注意
對
MultiSyncDataCollector和MultiaSyncDataCollector的支援是實驗性的，並且應始終首選SyncDataCollector。 如果需要在單個節點上執行多個並行的環境，請考慮使用ParallelEnv實例。collector_kwargs (dict 或 list, optional) – 要傳遞到遠端資料收集器的參數字典。 如果提供列表，則每個元素將對應於專用收集器的個別關鍵字引數集。
num_workers_per_collector (int, optional) – 要在遠端節點上使用的環境建構子副本數。 預設值為 1 (每個收集器一個環境)。 在單個工作節點上，所有子工作器都將執行相同的環境。 如果需要執行不同的環境，則應將它們分派到不同的工作節點，而不是子節點。
sync (bool, optional) – 如果
True，則產生的 tensordict 是在每個節點上收集的所有 tensordict 的堆疊。 如果False(預設)，則每個 tensordict 來自一個單獨的節點，以“先到先服務”的方式。slurm_kwargs (dict) – 要傳遞給 submitit 執行器的參數字典。
update_after_each_batch (bool, optional) – 若為
True，則權重將在每次資料收集後更新。對於sync=True，這表示所有 worker 都會看到其權重已更新。對於sync=False，則只會更新收集到資料的 worker。預設值為False，即更新必須通過update_policy_weights_()手動執行。max_weight_update_interval (int, optional) – 在 worker 的策略權重更新之前，可以收集的最大批次數量。對於同步資料收集，此參數會被
update_after_each_batch覆蓋。對於非同步資料收集，即使update_after_each_batch已開啟，也可能發生某個 worker 在一段時間內沒有看到其參數被更新的情況。預設值為 -1 (不強制更新)。launcher (str, optional) – 該如何啟動任務。可以是 “submitit” 或 “mp”（用於多重處理）之一。前者可以在多個節點上啟動任務，而後者僅在一台機器上啟動任務。“submitit” 需要安裝同名的函式庫。要了解更多關於 submitit 的資訊，請訪問 https://github.com/facebookincubator/submitit。預設值為 “submitit”。
tcp_port (int, optional) – 要使用的 TCP 連接埠。預設值為 10003。
visible_devices (list of Union[int, torch.device, str], optional) – 一個列表，其長度與節點數量相同，包含用於將資料傳遞給主節點的裝置。
tensorpipe_options (dict, optional) – 一個關鍵字引數字典，用於傳遞給
torch.distributed.rpc.TensorPipeRpcBackendOption。
