torchrl 0.11.0__cp314-cp314-macosx_11_0_arm64.whl

sota-implementations/iql/utils.py

@@ -0,0 +1,437 @@
+# 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.nn
+import torch.optim
+from tensordict.nn import InteractionType, TensorDictModule
+from tensordict.nn.distributions import NormalParamExtractor
+from torch.distributions import Categorical
+
+from torchrl.collectors import SyncDataCollector
+from torchrl.data import (
+    Composite,
+    LazyMemmapStorage,
+    TensorDictPrioritizedReplayBuffer,
+    TensorDictReplayBuffer,
+)
+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,
+    TransformedEnv,
+)
+
+from torchrl.envs.libs.gym import GymEnv, set_gym_backend
+from torchrl.envs.utils import ExplorationType, set_exploration_type
+from torchrl.modules import (
+    MLP,
+    ProbabilisticActor,
+    SafeModule,
+    TanhNormal,
+    ValueOperator,
+)
+from torchrl.objectives import DiscreteIQLLoss, HardUpdate, IQLLoss, SoftUpdate
+from torchrl.record import VideoRecorder
+from torchrl.trainers.helpers.models import ACTIVATIONS
+
+
+# ====================================================================
+# Environment utils
+# -----------------
+
+
+def env_maker(cfg, device="cpu", from_pixels=False):
+    lib = cfg.env.backend
+    if lib in ("gym", "gymnasium"):
+        with set_gym_backend(lib):
+            return GymEnv(
+                cfg.env.name,
+                device=device,
+                from_pixels=from_pixels,
+                pixels_only=False,
+                categorical_action_encoding=True,
+            )
+    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,
+):
+    transformed_env = TransformedEnv(
+        env,
+        Compose(
+            InitTracker(),
+            DoubleToFloat(),
+            RewardSum(),
+        ),
+    )
+    return transformed_env
+
+
+def make_environment(cfg, train_num_envs=1, eval_num_envs=1, logger=None):
+    """Make environments for training and evaluation."""
+    maker = functools.partial(env_maker, cfg)
+    parallel_env = ParallelEnv(
+        train_num_envs,
+        EnvCreator(maker),
+        serial_for_single=True,
+    )
+    parallel_env.set_seed(cfg.env.seed)
+
+    train_env = apply_env_transforms(parallel_env)
+
+    maker = functools.partial(env_maker, cfg, from_pixels=cfg.logger.video)
+    eval_env = TransformedEnv(
+        ParallelEnv(
+            eval_num_envs,
+            EnvCreator(maker),
+            serial_for_single=True,
+        ),
+        train_env.transform.clone(),
+    )
+    if cfg.logger.video:
+        eval_env.insert_transform(
+            0, VideoRecorder(logger, tag="rendered", in_keys=["pixels"])
+        )
+    return train_env, eval_env
+
+
+# ====================================================================
+# Collector and replay buffer
+# ---------------------------
+
+
+def make_collector(cfg, train_env, actor_model_explore, compile_mode):
+    """Make collector."""
+    device = cfg.collector.device
+    if device in ("", None):
+        if torch.cuda.is_available():
+            device = torch.device("cuda:0")
+        else:
+            device = torch.device("cpu")
+    collector = SyncDataCollector(
+        train_env,
+        actor_model_explore,
+        frames_per_batch=cfg.collector.frames_per_batch,
+        init_random_frames=cfg.collector.init_random_frames,
+        max_frames_per_traj=cfg.collector.max_frames_per_traj,
+        total_frames=cfg.collector.total_frames,
+        device=device,
+        compile_policy={"mode": compile_mode} if compile_mode else False,
+        cudagraph_policy={"warmup": 10} if cfg.compile.cudagraphs else False,
+    )
+    collector.set_seed(cfg.env.seed)
+    return collector
+
+
+def make_replay_buffer(
+    batch_size,
+    prb=False,
+    buffer_size=1000000,
+    scratch_dir=None,
+    device="cpu",
+    prefetch=3,
+):
+    if prb:
+        replay_buffer = TensorDictPrioritizedReplayBuffer(
+            alpha=0.7,
+            beta=0.5,
+            pin_memory=False,
+            prefetch=prefetch,
+            storage=LazyMemmapStorage(
+                buffer_size,
+                scratch_dir=scratch_dir,
+                device=device,
+            ),
+            batch_size=batch_size,
+        )
+    else:
+        replay_buffer = TensorDictReplayBuffer(
+            pin_memory=False,
+            prefetch=prefetch,
+            storage=LazyMemmapStorage(
+                buffer_size,
+                scratch_dir=scratch_dir,
+                device=device,
+            ),
+            batch_size=batch_size,
+        )
+    return replay_buffer
+
+
+def make_offline_replay_buffer(rb_cfg):
+    data = D4RLExperienceReplay(
+        dataset_id=rb_cfg.dataset,
+        split_trajs=False,
+        batch_size=rb_cfg.batch_size,
+        # We use drop_last to avoid recompiles (and dynamic shapes)
+        sampler=SamplerWithoutReplacement(drop_last=True),
+        prefetch=4,
+        direct_download=True,
+    )
+
+    data.append_transform(DoubleToFloat())
+
+    return data
+
+
+# ====================================================================
+# Model
+# -----
+#
+# We give one version of the model for learning from pixels, and one for state.
+# TorchRL comes in handy at this point, as the high-level interactions with
+# these models is unchanged, regardless of the modality.
+#
+
+
+def make_iql_model(cfg, train_env, eval_env, device="cpu"):
+    model_cfg = cfg.model
+
+    in_keys = ["observation"]
+    action_spec = train_env.action_spec_unbatched
+    actor_net, q_net, value_net = make_iql_modules_state(model_cfg, eval_env)
+
+    out_keys = ["loc", "scale"]
+
+    actor_module = TensorDictModule(actor_net, in_keys=in_keys, out_keys=out_keys)
+
+    # We use a ProbabilisticActor to make sure that we map the
+    # network output to the right space using a TanhDelta
+    # distribution.
+    actor = ProbabilisticActor(
+        module=actor_module,
+        in_keys=["loc", "scale"],
+        spec=action_spec,
+        distribution_class=TanhNormal,
+        distribution_kwargs={
+            "low": action_spec.space.low.to(device),
+            "high": action_spec.space.high.to(device),
+            "tanh_loc": False,
+        },
+        default_interaction_type=ExplorationType.RANDOM,
+    )
+
+    in_keys = ["observation", "action"]
+
+    out_keys = ["state_action_value"]
+    qvalue = ValueOperator(
+        in_keys=in_keys,
+        out_keys=out_keys,
+        module=q_net,
+    )
+    in_keys = ["observation"]
+    out_keys = ["state_value"]
+    value_net = ValueOperator(
+        in_keys=in_keys,
+        out_keys=out_keys,
+        module=value_net,
+    )
+    model = torch.nn.ModuleList([actor, qvalue, value_net]).to(device)
+    # init nets
+    with torch.no_grad(), set_exploration_type(ExplorationType.RANDOM):
+        td = eval_env.fake_tensordict()
+        td = td.to(device)
+        for net in model:
+            net(td)
+
+    return model
+
+
+def make_iql_modules_state(model_cfg, proof_environment):
+    action_spec = proof_environment.action_spec_unbatched
+
+    actor_net_kwargs = {
+        "num_cells": model_cfg.hidden_sizes,
+        "out_features": 2 * action_spec.shape[-1],
+        "activation_class": ACTIVATIONS[model_cfg.activation],
+    }
+    actor_net = MLP(**actor_net_kwargs)
+    actor_extractor = NormalParamExtractor(
+        scale_mapping=f"biased_softplus_{model_cfg.default_policy_scale}",
+        scale_lb=model_cfg.scale_lb,
+    )
+    actor_net = torch.nn.Sequential(actor_net, actor_extractor)
+
+    qvalue_net_kwargs = {
+        "num_cells": model_cfg.hidden_sizes,
+        "out_features": 1,
+        "activation_class": ACTIVATIONS[model_cfg.activation],
+    }
+
+    q_net = MLP(**qvalue_net_kwargs)
+
+    # Define Value Network
+    value_net_kwargs = {
+        "num_cells": model_cfg.hidden_sizes,
+        "out_features": 1,
+        "activation_class": ACTIVATIONS[model_cfg.activation],
+    }
+    value_net = MLP(**value_net_kwargs)
+
+    return actor_net, q_net, value_net
+
+
+def make_discrete_iql_model(cfg, train_env, eval_env, device):
+    """Make discrete IQL agent."""
+    # Define Actor Network
+    in_keys = ["observation"]
+    action_spec = train_env.action_spec_unbatched
+    # Define Actor Network
+    in_keys = ["observation"]
+
+    actor_net = MLP(
+        num_cells=cfg.model.hidden_sizes,
+        out_features=action_spec.space.n,
+        activation_class=ACTIVATIONS[cfg.model.activation],
+        device=device,
+    )
+
+    actor_module = SafeModule(
+        module=actor_net,
+        in_keys=in_keys,
+        out_keys=["logits"],
+    )
+    actor = ProbabilisticActor(
+        spec=Composite(action=eval_env.action_spec_unbatched).to(device),
+        module=actor_module,
+        in_keys=["logits"],
+        out_keys=["action"],
+        distribution_class=Categorical,
+        distribution_kwargs={},
+        default_interaction_type=InteractionType.RANDOM,
+        return_log_prob=False,
+    )
+
+    # Define Critic Network
+    qvalue_net = MLP(
+        num_cells=cfg.model.hidden_sizes,
+        out_features=action_spec.space.n,
+        activation_class=ACTIVATIONS[cfg.model.activation],
+        device=device,
+    )
+    qvalue = TensorDictModule(
+        in_keys=["observation"],
+        out_keys=["state_action_value"],
+        module=qvalue_net,
+    )
+
+    # Define Value Network
+    value_net = MLP(
+        num_cells=cfg.model.hidden_sizes,
+        out_features=1,
+        activation_class=ACTIVATIONS[cfg.model.activation],
+        device=device,
+    )
+    value_net = TensorDictModule(
+        in_keys=["observation"],
+        out_keys=["state_value"],
+        module=value_net,
+    )
+
+    model = torch.nn.ModuleList([actor, qvalue, value_net])
+    # init nets
+    with torch.no_grad(), set_exploration_type(ExplorationType.RANDOM):
+        td = eval_env.fake_tensordict()
+        td = td.to(device)
+        for net in model:
+            net(td)
+
+    return model
+
+
+# ====================================================================
+# IQL Loss
+# ---------
+
+
+def make_loss(loss_cfg, model, device):
+    loss_module = IQLLoss(
+        model[0],
+        model[1],
+        value_network=model[2],
+        loss_function=loss_cfg.loss_function,
+        temperature=loss_cfg.temperature,
+        expectile=loss_cfg.expectile,
+    )
+    loss_module.make_value_estimator(gamma=loss_cfg.gamma, device=device)
+    target_net_updater = SoftUpdate(loss_module, tau=loss_cfg.tau)
+
+    return loss_module, target_net_updater
+
+
+def make_discrete_loss(loss_cfg, model, device):
+    loss_module = DiscreteIQLLoss(
+        model[0],
+        model[1],
+        value_network=model[2],
+        loss_function=loss_cfg.loss_function,
+        temperature=loss_cfg.temperature,
+        expectile=loss_cfg.expectile,
+        action_space="categorical",
+    )
+    loss_module.make_value_estimator(gamma=loss_cfg.gamma, device=device)
+    target_net_updater = HardUpdate(
+        loss_module, value_network_update_interval=loss_cfg.hard_update_interval
+    )
+
+    return loss_module, target_net_updater
+
+
+def make_iql_optimizer(optim_cfg, loss_module):
+    critic_params = list(loss_module.qvalue_network_params.flatten_keys().values())
+    actor_params = list(loss_module.actor_network_params.flatten_keys().values())
+    value_params = list(loss_module.value_network_params.flatten_keys().values())
+
+    optimizer_actor = torch.optim.Adam(
+        actor_params,
+        lr=optim_cfg.lr,
+        weight_decay=optim_cfg.weight_decay,
+    )
+    optimizer_critic = torch.optim.Adam(
+        critic_params,
+        lr=optim_cfg.lr,
+        weight_decay=optim_cfg.weight_decay,
+    )
+    optimizer_value = torch.optim.Adam(
+        value_params,
+        lr=optim_cfg.lr,
+        weight_decay=optim_cfg.weight_decay,
+    )
+    return optimizer_actor, optimizer_critic, optimizer_value
+
+
+# ====================================================================
+# General utils
+# ---------
+
+
+def log_metrics(logger, metrics, step):
+    if logger is not None:
+        for metric_name, metric_value in metrics.items():
+            logger.log_scalar(metric_name, metric_value, step)
+
+
+def dump_video(module):
+    if isinstance(module, VideoRecorder):
+        module.dump()

sota-implementations/multiagent/README.md

@@ -0,0 +1,74 @@
+# Multi-agent examples
+
+In this folder we provide a set of multi-agent example scripts using the [VMAS](https://github.com/proroklab/VectorizedMultiAgentSimulator) simulator.
+
+<p align="center">
+<img src="https://pytorch.s3.amazonaws.com/torchrl/github-artifacts/img/marl_vmas.png" width="600px">
+</p>
+
+<center><i>The MARL algorithms contained in the scripts of this folder run on three multi-robot tasks in VMAS.</i></center>
+
+For more details on the experiment setup and the environments please refer to the corresponding section of the appendix in the [TorchRL paper](https://arxiv.org/abs/2306.00577).
+
+> [!NOTE]  
+> If you are interested in Multi-Agent Reinforcement Learning (MARL) in TorchRL, check out [BenchMARL](https://github.com/facebookresearch/BenchMARL):
+> a benchmarking library where you 
+> can train and compare MARL algorithms, tasks, and models using TorchRL!
+
+## Using the scripts
+
+### Install
+
+First you need to install vmas and the dependencies of the scripts.
+
+Install torchrl and tensordict following repo instructions.
+
+Install vmas and dependencies:
+
+```bash
+pip install vmas
+pip install wandb "moviepy<2.0.0"
+pip install hydra-core
+```
+
+### Run
+
+To run the scripts just execute the corresponding python file after having modified the corresponding config 
+according to your needs.
+The config can be found in the .yaml file with the same name.
+
+For example:
+```bash
+python mappo_ippo.py
+```
+
+You can even change the config from the command line like:
+
+```bash
+python mappo_ippo.py --m env.scenario_name=navigation
+```
+
+### Computational demand
+The scripts are set up for collecting many frames, if your compute is limited, you can change the "frames_per_batch" 
+and "num_epochs" parameters to reduce compute requirements.
+
+### Script structure
+
+The scripts are self-contained.
+This means that all the code you will need to look at is contained in the script file. 
+No helper functions are used.
+
+The structure of scripts follows this order:
+- Configuration dictionary for the script
+- Environment creation
+- Modules creation
+- Collector instantiation
+- Replay buffer instantiation
+- Loss module creation
+- Training loop (with inner minibatch loops)
+- Evaluation run (at the desired frequency)
+
+Logging is done by default to wandb.
+The logging backend can be changed in the config files to one of "wandb", "tensorboard", "csv", "mlflow".
+
+All the scripts follow the same on-policy training structure so that results can be compared across different algorithms.