RayCollector¶
- class torchrl.collectors.distributed.RayCollector(create_env_fn: ~typing.Union[~typing.Callable, ~torchrl.envs.common.EnvBase, ~typing.List[~typing.Callable], ~typing.List[~torchrl.envs.common.EnvBase]], policy: ~typing.Callable[[~tensordict._td.TensorDict], ~tensordict._td.TensorDict], *, frames_per_batch: int, total_frames: int = -1, device: ~typing.Optional[~typing.Union[~torch.device, ~typing.List[~torch.device]]] = None, storing_device: ~typing.Optional[~typing.Union[~torch.device, ~typing.List[~torch.device]]] = None, env_device: ~typing.Optional[~typing.Union[~torch.device, ~typing.List[~torch.device]]] = None, policy_device: ~typing.Optional[~typing.Union[~torch.device, ~typing.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: ~typing.Callable[[~tensordict._td.TensorDict], ~tensordict._td.TensorDict] = <class 'torchrl.collectors.collectors.SyncDataCollector'>, collector_kwargs: ~typing.Optional[~typing.Union[~typing.Dict, ~typing.List[~typing.Dict]]] = None, num_workers_per_collector: int = 1, sync: bool = False, ray_init_config: ~typing.Optional[~typing.Dict] = None, remote_configs: ~typing.Optional[~typing.Union[~typing.Dict, ~typing.List[~typing.Dict]]] = None, num_collectors: ~typing.Optional[int] = None, update_after_each_batch=False, max_weight_update_interval=-1)[source]¶
Distributed data collector with Ray backend.
This Python class serves as a ray-based solution to instantiate and coordinate multiple data collectors in a distributed cluster. Like TorchRL non-distributed collectors, this collector is an iterable that yields TensorDicts until a target number of collected frames is reached, but handles distributed data collection under the hood.
The class dictionary input parameter “ray_init_config” can be used to provide the kwargs to call Ray initialization method ray.init(). If “ray_init_config” is not provided, the default behavior is to autodetect an existing Ray cluster or start a new Ray instance locally if no existing cluster is found. Refer to Ray documentation for advanced initialization kwargs.
Similarly, dictionary input parameter “remote_configs” can be used to specify the kwargs for ray.remote() when called to create each remote collector actor, including collector compute resources.The sum of all collector resources should be available in the cluster. Refer to Ray documentation for advanced configuration of the ray.remote() method. Default kwargs are:
>>> kwargs = { ... "num_cpus": 1, ... "num_gpus": 0.2, ... "memory": 2 * 1024 ** 3, ... }
The coordination between collector instances can be specified as “synchronous” or “asynchronous”. In synchronous coordination, this class waits for all remote collectors to collect a rollout, concatenates all rollouts into a single TensorDict instance and finally yields the concatenated data. On the other hand, if the coordination is to be carried out asynchronously, this class provides the rollouts as they become available from individual remote collectors.
- Parameters:
create_env_fn (Callable or List[Callabled]) – list of Callables, each returning an instance of
EnvBase
.policy (Callable) –
Policy to be executed in the environment. Must accept
tensordict.tensordict.TensorDictBase
object as input. IfNone
is provided, the policy used will be aRandomPolicy
instance with the environmentaction_spec
. Accepted policies are usually subclasses ofTensorDictModuleBase
. This is the recommended usage of the collector. Other callables are accepted too: If the policy is not aTensorDictModuleBase
(e.g., a regularModule
instances) it will be wrapped in a nn.Module first. Then, the collector will try to assess if these modules require wrapping in aTensorDictModule
or not. - If the policy forward signature matches any offorward(self, tensordict)
,forward(self, td)
orforward(self, <anything>: TensorDictBase)
(or any typing with a single argument typed as a subclass ofTensorDictBase
) then the policy won’t be wrapped in aTensorDictModule
.In all other cases an attempt to wrap it will be undergone as such:
TensorDictModule(policy, in_keys=env_obs_key, out_keys=env.action_keys)
.
- Keyword Arguments:
frames_per_batch (int) – A keyword-only argument representing the total number of elements in a batch.
total_frames (int, Optional) – lower bound of the total number of frames returned by the collector. The iterator will stop once the total number of frames equates or exceeds the total number of frames passed to the collector. Default value is -1, which mean no target total number of frames (i.e. the collector will run indefinitely).
device (int, str or torch.device, optional) – The generic device of the collector. The
device
args fills any non-specified device: ifdevice
is notNone
and any ofstoring_device
,policy_device
orenv_device
is not specified, its value will be set todevice
. Defaults toNone
(No default device). Lists of devices are supported.storing_device (int, str or torch.device, optional) – The remote device on which the output
TensorDict
will be stored. Ifdevice
is passed andstoring_device
isNone
, it will default to the value indicated bydevice
. For long trajectories, it may be necessary to store the data on a different device than the one where the policy and env are executed. Defaults toNone
(the output tensordict isn’t on a specific device, leaf tensors sit on the device where they were created). Lists of devices are supported.env_device (int, str or torch.device, optional) – The remote device on which the environment should be cast (or executed if that functionality is supported). If not specified and the env has a non-
None
device,env_device
will default to that value. Ifdevice
is passed andenv_device=None
, it will default todevice
. If the value as such specified ofenv_device
differs frompolicy_device
and one of them is notNone
, the data will be cast toenv_device
before being passed to the env (i.e., passing different devices to policy and env is supported). Defaults toNone
. Lists of devices are supported.policy_device (int, str or torch.device, optional) – The remote device on which the policy should be cast. If
device
is passed andpolicy_device=None
, it will default todevice
. If the value as such specified ofpolicy_device
differs fromenv_device
and one of them is notNone
, the data will be cast topolicy_device
before being passed to the policy (i.e., passing different devices to policy and env is supported). Defaults toNone
. Lists of devices are supported.create_env_kwargs (dict, optional) – Dictionary of kwargs for
create_env_fn
.max_frames_per_traj (int, optional) – Maximum steps per trajectory. Note that a trajectory can span across multiple batches (unless
reset_at_each_iter
is set toTrue
, see below). Once a trajectory reachesn_steps
, the environment is reset. If the environment wraps multiple environments together, the number of steps is tracked for each environment independently. Negative values are allowed, in which case this argument is ignored. Defaults toNone
(i.e., no maximum number of steps).init_random_frames (int, optional) – Number of frames for which the policy is ignored before it is called. This feature is mainly intended to be used in offline/model-based settings, where a batch of random trajectories can be used to initialize training. If provided, it will be rounded up to the closest multiple of frames_per_batch. Defaults to
None
(i.e. no random frames).reset_at_each_iter (bool, optional) – Whether environments should be reset at the beginning of a batch collection. Defaults to
False
.postproc (Callable, optional) – A post-processing transform, such as a
Transform
or aMultiStep
instance. Defaults toNone
.split_trajs (bool, optional) – Boolean indicating whether the resulting TensorDict should be split according to the trajectories. See
split_trajectories()
for more information. Defaults toFalse
.exploration_type (ExplorationType, optional) – interaction mode to be used when collecting data. Must be one of
torchrl.envs.utils.ExplorationType.DETERMINISTIC
,torchrl.envs.utils.ExplorationType.RANDOM
,torchrl.envs.utils.ExplorationType.MODE
ortorchrl.envs.utils.ExplorationType.MEAN
.collector_class (Python class) – a collector class to be remotely instantiated. Can be
SyncDataCollector
,MultiSyncDataCollector
,MultiaSyncDataCollector
or a derived class of these. Defaults toSyncDataCollector
.collector_kwargs (dict or list, optional) – a dictionary of parameters to be passed to the remote data-collector. If a list is provided, each element will correspond to an individual set of keyword arguments for the dedicated collector.
num_workers_per_collector (int) – the number of copies of the env constructor that is to be used on the remote nodes. Defaults to 1 (a single env per collector). On a single worker node all the sub-workers will be executing the same environment. If different environments need to be executed, they should be dispatched across worker nodes, not subnodes.
ray_init_config (dict, Optional) – kwargs used to call ray.init().
remote_configs (list of dicts, Optional) – ray resource specs for each remote collector. A single dict can be provided as well, and will be used in all collectors.
num_collectors (int, Optional) – total number of collectors to be instantiated.
sync (bool) – if
True
, the resulting tensordict is a stack of all the tensordicts collected on each node. IfFalse
(default), each tensordict results from a separate node in a “first-ready, first-served” fashion.update_after_each_batch (bool, optional) – if
True
, the weights will be updated after each collection. Forsync=True
, this means that all workers will see their weights updated. Forsync=False
, only the worker from which the data has been gathered will be updated. Defaults toFalse
, i.e. updates have to be executed manually throughtorchrl.collectors.distributed.RayDistributedCollector.update_policy_weights_()
max_weight_update_interval (int, optional) – the maximum number of batches that can be collected before the policy weights of a worker is updated. For sync collections, this parameter is overwritten by
update_after_each_batch
. For async collections, it may be that one worker has not seen its parameters being updated for a certain time even ifupdate_after_each_batch
is turned on. Defaults to -1 (no forced update).
Examples
>>> 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
- add_collectors(create_env_fn, num_envs, policy, collector_kwargs, remote_configs)[source]¶
Creates and adds a number of remote collectors to the set.
- load_state_dict(state_dict: Union[OrderedDict, List[OrderedDict]]) None [source]¶
Calls parent method for each remote collector.
- set_seed(seed: int, static_seed: bool = False) List[int] [source]¶
Calls parent method for each remote collector iteratively and returns final seed.