Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Doc] Indicate necessary context to run multiprocessed collectors in doc #2126

Merged
merged 8 commits into from
Apr 30, 2024
Merged

[Doc] Indicate necessary context to run multiprocessed collectors in doc #2126

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
38fee42
Updated Multisynccollector
Apr 26, 2024
6cbc738
Update multiasync docstring.
GJBoth Apr 26, 2024
f9960af
Code block to inline.
GJBoth Apr 26, 2024
ba10554
Update torchrl/collectors/collectors.py
GJBoth Apr 29, 2024
1fae6d2
Update torchrl/collectors/collectors.py
GJBoth Apr 29, 2024
1029f11
Update sync and indent shutdown.
GJBoth Apr 29, 2024
773618d
Merge remote-tracking branch 'origin/main' into update_collector_docs
vmoens Apr 29, 2024
d425c25
amend
vmoens Apr 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 61 additions & 41 deletions torchrl/collectors/collectors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1884,28 +1884,40 @@ class MultiSyncDataCollector(_MultiDataCollector):
trajectory and the start of the next collection.
This class can be safely used with online RL sota-implementations.

.. note:: Python requires multiprocessed code to be instantiated within a main guard:

>>> from torchrl.collectors import MultiSyncDataCollector
>>> if __name__ == "__main__":
... # Create your collector here

See https://docs.python.org/3/library/multiprocessing.html for more info.

Examples:
>>> from torchrl.envs.libs.gym import GymEnv
>>> from torchrl.envs import StepCounter
>>> from tensordict.nn import TensorDictModule
>>> from torch import nn
>>> env_maker = lambda: TransformedEnv(GymEnv("Pendulum-v1", device="cpu"), StepCounter(max_steps=50))
>>> policy = TensorDictModule(nn.Linear(3, 1), in_keys=["observation"], out_keys=["action"])
>>> collector = MultiSyncDataCollector(
... create_env_fn=[env_maker, env_maker],
... policy=policy,
... total_frames=2000,
... max_frames_per_traj=50,
... frames_per_batch=200,
... init_random_frames=-1,
... reset_at_each_iter=False,
... devices="cpu",
... storing_devices="cpu",
... )
>>> for i, data in enumerate(collector):
... if i == 2:
... print(data)
... break
>>> from torchrl.collectors import MultiSyncDataCollector
>>> if __name__ == "__main__":
... env_maker = lambda: GymEnv("Pendulum-v1", device="cpu")
... policy = TensorDictModule(nn.Linear(3, 1), in_keys=["observation"], out_keys=["action"])
... collector = MultiSyncDataCollector(
... create_env_fn=[env_maker, env_maker],
... policy=policy,
... total_frames=2000,
... max_frames_per_traj=50,
... frames_per_batch=200,
... init_random_frames=-1,
... reset_at_each_iter=False,
... device="cpu",
... storing_device="cpu",
... cat_results="stack",
... )
... for i, data in enumerate(collector):
... if i == 2:
... print(data)
... break
... collector.shutdown()
... del collector
TensorDict(
fields={
action: Tensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.float32, is_shared=False),
Expand All @@ -1932,8 +1944,6 @@ class MultiSyncDataCollector(_MultiDataCollector):
batch_size=torch.Size([200]),
device=cpu,
is_shared=False)
>>> collector.shutdown()
>>> del collector

"""

Expand Down Expand Up @@ -1987,7 +1997,6 @@ def _queue_len(self) -> int:
return self.num_workers

def iterator(self) -> Iterator[TensorDictBase]:

cat_results = self.cat_results
if cat_results is None:
cat_results = 0
Expand Down Expand Up @@ -2232,27 +2241,40 @@ class MultiaSyncDataCollector(_MultiDataCollector):
the batch of rollouts is collected and the next call to the iterator.
This class can be safely used with offline RL sota-implementations.

Examples:
.. note:: Python requires multiprocessed code to be instantiated within a main guard:

>>> from torchrl.collectors import MultiaSyncDataCollector
>>> if __name__ == "__main__":
... # Create your collector here

See https://docs.python.org/3/library/multiprocessing.html for more info.

Examples:
>>> from torchrl.envs.libs.gym import GymEnv
>>> from tensordict.nn import TensorDictModule
>>> from torch import nn
>>> env_maker = lambda: GymEnv("Pendulum-v1", device="cpu")
>>> policy = TensorDictModule(nn.Linear(3, 1), in_keys=["observation"], out_keys=["action"])
>>> collector = MultiaSyncDataCollector(
... create_env_fn=[env_maker, env_maker],
... policy=policy,
... total_frames=2000,
... max_frames_per_traj=50,
... frames_per_batch=200,
... init_random_frames=-1,
... reset_at_each_iter=False,
... devices="cpu",
... storing_devices="cpu",
... )
>>> for i, data in enumerate(collector):
... if i == 2:
... print(data)
... break
>>> from torchrl.collectors import MultiaSyncDataCollector
>>> if __name__ == "__main__":
... env_maker = lambda: GymEnv("Pendulum-v1", device="cpu")
... policy = TensorDictModule(nn.Linear(3, 1), in_keys=["observation"], out_keys=["action"])
... collector = MultiaSyncDataCollector(
... create_env_fn=[env_maker, env_maker],
... policy=policy,
... total_frames=2000,
... max_frames_per_traj=50,
... frames_per_batch=200,
... init_random_frames=-1,
... reset_at_each_iter=False,
... device="cpu",
... storing_device="cpu",
... cat_results="stack",
... )
... for i, data in enumerate(collector):
... if i == 2:
... print(data)
... break
... collector.shutdown()
... del collector
TensorDict(
fields={
action: Tensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.float32, is_shared=False),
Expand All @@ -2279,8 +2301,6 @@ class MultiaSyncDataCollector(_MultiDataCollector):
batch_size=torch.Size([200]),
device=cpu,
is_shared=False)
>>> collector.shutdown()
>>> del collector

"""

Expand Down
10 changes: 8 additions & 2 deletions torchrl/data/tensor_specs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1597,12 +1597,18 @@ def __init__(
if high is not None:
raise TypeError(self.CONFLICTING_KWARGS.format("high", "maximum"))
high = kwargs.pop("maximum")
warnings.warn("Maximum is deprecated since v0.4.0, using high instead.", category=DeprecationWarning)
warnings.warn(
"Maximum is deprecated since v0.4.0, using high instead.",
category=DeprecationWarning,
)
if "minimum" in kwargs:
if low is not None:
raise TypeError(self.CONFLICTING_KWARGS.format("low", "minimum"))
low = kwargs.pop("minimum")
warnings.warn("Minimum is deprecated since v0.4.0, using low instead.", category=DeprecationWarning)
warnings.warn(
"Minimum is deprecated since v0.4.0, using low instead.",
category=DeprecationWarning,
)
domain = kwargs.pop("domain", "continuous")
if len(kwargs):
raise TypeError(f"Got unrecognised kwargs {tuple(kwargs.keys())}.")
Expand Down
Loading