torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

sota-implementations/td3_bc/utils.py

@@ -0,0 +1,251 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import functools
+
+import torch
+from tensordict.nn import TensorDictModule, TensorDictSequential
+
+from torch import nn, optim
+from torchrl.data.datasets.d4rl import D4RLExperienceReplay
+from torchrl.data.replay_buffers import SamplerWithoutReplacement
+from torchrl.envs import (
+    CatTensors,
+    Compose,
+    DMControlEnv,
+    DoubleToFloat,
+    EnvCreator,
+    InitTracker,
+    ParallelEnv,
+    RewardSum,
+    StepCounter,
+    TransformedEnv,
+)
+from torchrl.envs.libs.gym import GymEnv, set_gym_backend
+from torchrl.envs.utils import ExplorationType, set_exploration_type
+from torchrl.modules import AdditiveGaussianModule, MLP, TanhModule, ValueOperator
+
+from torchrl.objectives import SoftUpdate
+from torchrl.objectives.td3_bc import TD3BCLoss
+from torchrl.record import VideoRecorder
+
+
+# ====================================================================
+# Environment utils
+# -----------------
+
+
+def env_maker(cfg, device="cpu", from_pixels=False):
+    lib = cfg.env.library
+    if lib in ("gym", "gymnasium"):
+        with set_gym_backend(lib):
+            return GymEnv(
+                cfg.env.name,
+                device=device,
+                from_pixels=from_pixels,
+                pixels_only=False,
+            )
+    elif lib == "dm_control":
+        env = DMControlEnv(
+            cfg.env.name, cfg.env.task, from_pixels=from_pixels, pixels_only=False
+        )
+        return TransformedEnv(
+            env, CatTensors(in_keys=env.observation_spec.keys(), out_key="observation")
+        )
+    else:
+        raise NotImplementedError(f"Unknown lib {lib}.")
+
+
+def apply_env_transforms(env, max_episode_steps):
+    transformed_env = TransformedEnv(
+        env,
+        Compose(
+            StepCounter(max_steps=max_episode_steps),
+            InitTracker(),
+            DoubleToFloat(),
+            RewardSum(),
+        ),
+    )
+    return transformed_env
+
+
+def make_environment(cfg, logger=None):
+    """Make environments for training and evaluation."""
+    partial = functools.partial(env_maker, cfg=cfg)
+    parallel_env = ParallelEnv(
+        cfg.logger.eval_envs,
+        EnvCreator(partial),
+        serial_for_single=True,
+    )
+    parallel_env.set_seed(cfg.env.seed)
+
+    train_env = apply_env_transforms(parallel_env, cfg.env.max_episode_steps)
+    return train_env
+
+
+# ====================================================================
+# Replay buffer
+# ---------------------------
+
+
+def make_offline_replay_buffer(rb_cfg, device):
+    data = D4RLExperienceReplay(
+        dataset_id=rb_cfg.dataset,
+        split_trajs=False,
+        batch_size=rb_cfg.batch_size,
+        # drop_last for compile
+        sampler=SamplerWithoutReplacement(drop_last=True),
+        prefetch=4,
+        direct_download=True,
+    )
+
+    data.append_transform(DoubleToFloat())
+    data.append_transform(lambda td: td.to(device))
+
+    return data
+
+
+# ====================================================================
+# Model
+# -----
+
+
+def make_td3_agent(cfg, train_env, device):
+    """Make TD3 agent."""
+    # Define Actor Network
+    in_keys = ["observation"]
+    action_spec = train_env.action_spec_unbatched.to(device)
+
+    actor_net = MLP(
+        num_cells=cfg.network.hidden_sizes,
+        out_features=action_spec.shape[-1],
+        activation_class=get_activation(cfg),
+        device=device,
+    )
+
+    in_keys_actor = in_keys
+    actor_module = TensorDictModule(
+        actor_net,
+        in_keys=in_keys_actor,
+        out_keys=["param"],
+    )
+    actor = TensorDictSequential(
+        actor_module,
+        TanhModule(
+            in_keys=["param"],
+            out_keys=["action"],
+            spec=action_spec,
+        ),
+    )
+
+    # Define Critic Network
+    qvalue_net = MLP(
+        num_cells=cfg.network.hidden_sizes,
+        out_features=1,
+        activation_class=get_activation(cfg),
+        device=device,
+    )
+
+    qvalue = ValueOperator(
+        in_keys=["action"] + in_keys,
+        module=qvalue_net,
+    )
+
+    model = nn.ModuleList([actor, qvalue])
+
+    # init nets
+    with torch.no_grad(), set_exploration_type(ExplorationType.RANDOM):
+        td = train_env.fake_tensordict()
+        td = td.to(device)
+        for net in model:
+            net(td)
+    del td
+
+    # Exploration wrappers:
+    actor_model_explore = TensorDictSequential(
+        model[0],
+        AdditiveGaussianModule(
+            sigma_init=1,
+            sigma_end=1,
+            mean=0,
+            std=0.1,
+            spec=action_spec,
+            device=device,
+        ),
+    )
+    return model, actor_model_explore
+
+
+# ====================================================================
+# TD3 Loss
+# ---------
+
+
+def make_loss_module(cfg, model):
+    """Make loss module and target network updater."""
+    # Create TD3 loss
+    loss_module = TD3BCLoss(
+        actor_network=model[0],
+        qvalue_network=model[1],
+        num_qvalue_nets=2,
+        loss_function=cfg.loss_function,
+        delay_actor=True,
+        delay_qvalue=True,
+        action_spec=model[0][1].spec,
+        policy_noise=cfg.policy_noise,
+        noise_clip=cfg.noise_clip,
+        alpha=cfg.alpha,
+    )
+    loss_module.make_value_estimator(gamma=cfg.gamma)
+
+    # Define Target Network Updater
+    target_net_updater = SoftUpdate(loss_module, eps=cfg.target_update_polyak)
+    return loss_module, target_net_updater
+
+
+def make_optimizer(cfg, loss_module):
+    critic_params = list(loss_module.qvalue_network_params.values(True, True))
+    actor_params = list(loss_module.actor_network_params.values(True, True))
+
+    optimizer_actor = optim.Adam(
+        actor_params,
+        lr=cfg.lr,
+        weight_decay=cfg.weight_decay,
+        eps=cfg.adam_eps,
+    )
+    optimizer_critic = optim.Adam(
+        critic_params,
+        lr=cfg.lr,
+        weight_decay=cfg.weight_decay,
+        eps=cfg.adam_eps,
+    )
+    return optimizer_actor, optimizer_critic
+
+
+# ====================================================================
+# General utils
+# ---------
+
+
+def log_metrics(logger, metrics, step):
+    for metric_name, metric_value in metrics.items():
+        logger.log_scalar(metric_name, metric_value, step)
+
+
+def get_activation(cfg):
+    if cfg.network.activation == "relu":
+        return nn.ReLU
+    elif cfg.network.activation == "tanh":
+        return nn.Tanh
+    elif cfg.network.activation == "leaky_relu":
+        return nn.LeakyReLU
+    else:
+        raise NotImplementedError
+
+
+def dump_video(module):
+    if isinstance(module, VideoRecorder):
+        module.dump()

torchrl/__init__.py

@@ -0,0 +1,144 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import warnings
+import weakref
+from warnings import warn
+
+import torch
+
+# Silence noisy dependency warning triggered at import time on older torch stacks.
+# (Emitted by tensordict when registering pytree nodes.)
+warnings.filterwarnings(
+    "ignore",
+    category=UserWarning,
+    message=r"torch\.utils\._pytree\._register_pytree_node is deprecated\.",
+)
+
+from tensordict import set_lazy_legacy  # noqa: E402
+
+from torch import multiprocessing as mp  # noqa: E402
+from torch.distributions.transforms import (  # noqa: E402
+    _InverseTransform,
+    ComposeTransform,
+)
+
+torch._C._log_api_usage_once("torchrl")
+
+set_lazy_legacy(False).set()
+
+from ._extension import _init_extension  # noqa: E402
+
+__version__ = None  # type: ignore
+try:
+    try:
+        from importlib.metadata import version as _dist_version
+    except ImportError:  # pragma: no cover
+        from importlib_metadata import version as _dist_version  # type: ignore
+
+    __version__ = _dist_version("torchrl")
+except Exception:
+    try:
+        from ._version import __version__
+    except Exception:
+        try:
+            from .version import __version__
+        except Exception:
+            __version__ = None  # type: ignore
+
+try:
+    from torch.compiler import is_dynamo_compiling
+except ImportError:
+    from torch._dynamo import is_compiling as is_dynamo_compiling
+
+_init_extension()
+
+from torchrl._utils import (  # noqa: E402
+    _get_default_mp_start_method,
+    auto_unwrap_transformed_env,
+    compile_with_warmup,
+    get_ray_default_runtime_env,
+    implement_for,
+    logger,
+    merge_ray_runtime_env,
+    set_auto_unwrap_transformed_env,
+    set_profiling_enabled,
+    timeit,
+)
+
+logger = logger
+
+# TorchRL's multiprocessing default.
+_preferred_start_method = _get_default_mp_start_method()
+if _preferred_start_method == "spawn":
+    try:
+        mp.set_start_method("spawn")
+    except RuntimeError as err:
+        if str(err).startswith("context has already been set"):
+            mp_start_method = mp.get_start_method()
+            if mp_start_method != "spawn":
+                warn(
+                    f"failed to set start method to spawn, "
+                    f"and current start method for mp is {mp_start_method}."
+                )
+
+# Filter warnings in subprocesses: True by default given the multiple optional
+# deps of the library. This can be turned on via `torchrl.filter_warnings_subprocess = False`.
+filter_warnings_subprocess = True
+
+_THREAD_POOL_INIT = torch.get_num_threads()
+
+
+# monkey-patch dist transforms until https://github.com/pytorch/pytorch/pull/135001/ finds a home
+@property
+def _inv(self):
+    """Patched version of Transform.inv.
+
+    Returns the inverse :class:`Transform` of this transform.
+
+    This should satisfy ``t.inv.inv is t``.
+    """
+    inv = None
+    if self._inv is not None:
+        inv = self._inv()
+    if inv is None:
+        inv = _InverseTransform(self)
+        if not is_dynamo_compiling():
+            self._inv = weakref.ref(inv)
+    return inv
+
+
+torch.distributions.transforms.Transform.inv = _inv
+
+
+@property
+def _inv(self):
+    inv = None
+    if self._inv is not None:
+        inv = self._inv()
+    if inv is None:
+        inv = ComposeTransform([p.inv for p in reversed(self.parts)])
+        if not is_dynamo_compiling():
+            self._inv = weakref.ref(inv)
+            inv._inv = weakref.ref(self)
+        else:
+            # We need inv.inv to be equal to self, but weakref can cause a graph break
+            inv._inv = lambda out=self: out
+
+    return inv
+
+
+ComposeTransform.inv = _inv
+
+__all__ = [
+    "auto_unwrap_transformed_env",
+    "compile_with_warmup",
+    "get_ray_default_runtime_env",
+    "implement_for",
+    "merge_ray_runtime_env",
+    "set_auto_unwrap_transformed_env",
+    "timeit",
+    "logger",
+    "logger",
+]

torchrl/_extension.py

@@ -0,0 +1,74 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import importlib.util
+import warnings
+
+from packaging.version import parse
+
+__version__ = None  # type: ignore
+try:
+    try:
+        from importlib.metadata import version as _dist_version
+    except ImportError:  # pragma: no cover
+        from importlib_metadata import version as _dist_version  # type: ignore
+
+    __version__ = _dist_version("torchrl")
+except Exception:
+    __version__ = None  # type: ignore
+
+try:
+    from .version import pytorch_version
+except ImportError:
+    pytorch_version = "unknown"
+
+
+def is_module_available(*modules: str) -> bool:
+    """Returns if a top-level module with :attr:`name` exists *without** importing it.
+
+    This is generally safer than try-catch block around a
+    `import X`. It avoids third party libraries breaking assumptions of some of
+    our tests, e.g., setting multiprocessing start method when imported
+    (see librosa/#747, torchvision/#544).
+    """
+    return all(importlib.util.find_spec(m) is not None for m in modules)
+
+
+def _init_extension():
+    if not is_module_available("torchrl._torchrl"):
+        warnings.warn("torchrl C++ extension is not available.")
+        return
+
+
+def _is_nightly(version):
+    if version is None:
+        return True
+    parsed_version = parse(version)
+    return parsed_version.local is not None
+
+
+if _is_nightly(__version__):
+    EXTENSION_WARNING = (
+        "Failed to import torchrl C++ binaries. Some modules (eg, prioritized replay buffers) may not work with your installation. "
+        "You seem to be using the nightly version of TorchRL. If this is a local install, there might be an issue with "
+        "the local installation. Here are some tips to debug this:\n"
+        " - make sure ninja and cmake were installed\n"
+        " - make sure you ran `python setup.py clean && python setup.py develop` and that no error was raised\n"
+        " - make sure the version of PyTorch you are using matches the one that was present in your virtual env during "
+        f"setup. This package was built with PyTorch {pytorch_version}. You can deactivate this warning by setting the environment variable `RL_WARNINGS=0`."
+    )
+
+else:
+    EXTENSION_WARNING = (
+        "Failed to import torchrl C++ binaries. Some modules (eg, prioritized replay buffers) may not work with your installation. "
+        "This is likely due to a discrepancy between your package version and the PyTorch version. "
+        "TorchRL does not tightly pin PyTorch versions to give users freedom, but the trade-off is that C++ extensions like "
+        "prioritized replay buffers can only be used with the PyTorch version they were built against. "
+        f"This package was built with PyTorch {pytorch_version}. "
+        "Workarounds include: (1) upgrading/downgrading PyTorch or TorchRL to compatible versions, "
+        "or (2) making a local install using `pip install git+https://github.com/pytorch/rl.git@<version>`. "
+        "You can deactivate this warning by setting the environment variable `RL_WARNINGS=0`."
+    )