torchrl 0.11.0__cp314-cp314t-win_amd64.whl

sota-implementations/multiagent/sac.py

@@ -0,0 +1,337 @@
+# 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 time
+
+import hydra
+import torch
+from tensordict.nn import TensorDictModule
+from tensordict.nn.distributions import NormalParamExtractor
+from torch import nn
+from torch.distributions import Categorical, OneHotCategorical
+from torchrl._utils import logger as torchrl_logger
+from torchrl.collectors import SyncDataCollector
+from torchrl.data import TensorDictReplayBuffer
+from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
+from torchrl.data.replay_buffers.storages import LazyTensorStorage
+from torchrl.envs import RewardSum, TransformedEnv
+from torchrl.envs.libs.vmas import VmasEnv
+from torchrl.envs.utils import ExplorationType, set_exploration_type
+from torchrl.modules import ProbabilisticActor, TanhNormal, ValueOperator
+from torchrl.modules.models.multiagent import MultiAgentMLP
+from torchrl.objectives import DiscreteSACLoss, SACLoss, SoftUpdate, ValueEstimators
+from utils.logging import init_logging, log_evaluation, log_training
+from utils.utils import DoneTransform
+
+
+def rendering_callback(env, td):
+    env.frames.append(env.render(mode="rgb_array", agent_index_focus=None))
+
+
+@hydra.main(version_base="1.1", config_path="", config_name="sac")
+def train(cfg: DictConfig):  # noqa: F821
+    # Device
+    cfg.train.device = "cpu" if not torch.cuda.device_count() else "cuda:0"
+    cfg.env.device = cfg.train.device
+
+    # Seeding
+    torch.manual_seed(cfg.seed)
+
+    # Sampling
+    cfg.env.vmas_envs = cfg.collector.frames_per_batch // cfg.env.max_steps
+    cfg.collector.total_frames = cfg.collector.frames_per_batch * cfg.collector.n_iters
+    cfg.buffer.memory_size = cfg.collector.frames_per_batch
+
+    # Create env and env_test
+    env = VmasEnv(
+        scenario=cfg.env.scenario_name,
+        num_envs=cfg.env.vmas_envs,
+        continuous_actions=cfg.env.continuous_actions,
+        max_steps=cfg.env.max_steps,
+        device=cfg.env.device,
+        seed=cfg.seed,
+        categorical_actions=cfg.env.categorical_actions,
+        # Scenario kwargs
+        **cfg.env.scenario,
+    )
+    env = TransformedEnv(
+        env,
+        RewardSum(in_keys=[env.reward_key], out_keys=[("agents", "episode_reward")]),
+    )
+
+    env_test = VmasEnv(
+        scenario=cfg.env.scenario_name,
+        num_envs=cfg.eval.evaluation_episodes,
+        continuous_actions=cfg.env.continuous_actions,
+        max_steps=cfg.env.max_steps,
+        device=cfg.env.device,
+        seed=cfg.seed,
+        # Scenario kwargs
+        **cfg.env.scenario,
+    )
+
+    # Policy
+    if cfg.env.continuous_actions:
+        actor_net = nn.Sequential(
+            MultiAgentMLP(
+                n_agent_inputs=env.full_observation_spec_unbatched[
+                    "agents", "observation"
+                ].shape[-1],
+                n_agent_outputs=2
+                * env.full_action_spec_unbatched["agents", "action"].shape[-1],
+                n_agents=env.n_agents,
+                centralised=False,
+                share_params=cfg.model.shared_parameters,
+                device=cfg.train.device,
+                depth=2,
+                num_cells=256,
+                activation_class=nn.Tanh,
+            ),
+            NormalParamExtractor(),
+        )
+        policy_module = TensorDictModule(
+            actor_net,
+            in_keys=[("agents", "observation")],
+            out_keys=[("agents", "loc"), ("agents", "scale")],
+        )
+
+        policy = ProbabilisticActor(
+            module=policy_module,
+            spec=env.full_action_spec_unbatched,
+            in_keys=[("agents", "loc"), ("agents", "scale")],
+            out_keys=[env.action_key],
+            distribution_class=TanhNormal,
+            distribution_kwargs={
+                "low": env.full_action_spec_unbatched[("agents", "action")].space.low,
+                "high": env.full_action_spec_unbatched[("agents", "action")].space.high,
+            },
+            return_log_prob=True,
+        )
+
+        # Critic
+        module = MultiAgentMLP(
+            n_agent_inputs=env.observation_spec["agents", "observation"].shape[-1]
+            + env.full_action_spec_unbatched["agents", "action"].shape[
+                -1
+            ],  # Q critic takes action and value
+            n_agent_outputs=1,
+            n_agents=env.n_agents,
+            centralised=cfg.model.centralised_critic,
+            share_params=cfg.model.shared_parameters,
+            device=cfg.train.device,
+            depth=2,
+            num_cells=256,
+            activation_class=nn.Tanh,
+        )
+        value_module = ValueOperator(
+            module=module,
+            in_keys=[("agents", "observation"), env.action_key],
+            out_keys=[("agents", "state_action_value")],
+        )
+    else:
+        actor_net = nn.Sequential(
+            MultiAgentMLP(
+                n_agent_inputs=env.observation_spec["agents", "observation"].shape[-1],
+                n_agent_outputs=env.full_action_spec_unbatched[
+                    "agents", "action"
+                ].space.n,
+                n_agents=env.n_agents,
+                centralised=False,
+                share_params=cfg.model.shared_parameters,
+                device=cfg.train.device,
+                depth=2,
+                num_cells=256,
+                activation_class=nn.Tanh,
+            ),
+        )
+        policy_module = TensorDictModule(
+            actor_net,
+            in_keys=[("agents", "observation")],
+            out_keys=[("agents", "logits")],
+        )
+        policy = ProbabilisticActor(
+            module=policy_module,
+            spec=env.full_action_spec_unbatched["agents", "action"],
+            in_keys=[("agents", "logits")],
+            out_keys=[env.action_key],
+            distribution_class=OneHotCategorical
+            if not cfg.env.categorical_actions
+            else Categorical,
+            return_log_prob=True,
+        )
+
+        # Critic
+        module = MultiAgentMLP(
+            n_agent_inputs=env.observation_spec["agents", "observation"].shape[-1],
+            n_agent_outputs=env.full_action_spec_unbatched["agents", "action"].space.n,
+            n_agents=env.n_agents,
+            centralised=cfg.model.centralised_critic,
+            share_params=cfg.model.shared_parameters,
+            device=cfg.train.device,
+            depth=2,
+            num_cells=256,
+            activation_class=nn.Tanh,
+        )
+        value_module = ValueOperator(
+            module=module,
+            in_keys=[("agents", "observation")],
+            out_keys=[("agents", "action_value")],
+        )
+
+    collector = SyncDataCollector(
+        env,
+        policy,
+        device=cfg.env.device,
+        storing_device=cfg.train.device,
+        frames_per_batch=cfg.collector.frames_per_batch,
+        total_frames=cfg.collector.total_frames,
+        postproc=DoneTransform(reward_key=env.reward_key, done_keys=env.done_keys),
+    )
+
+    replay_buffer = TensorDictReplayBuffer(
+        storage=LazyTensorStorage(cfg.buffer.memory_size, device=cfg.train.device),
+        sampler=SamplerWithoutReplacement(),
+        batch_size=cfg.train.minibatch_size,
+    )
+
+    if cfg.env.continuous_actions:
+        loss_module = SACLoss(
+            actor_network=policy,
+            qvalue_network=value_module,
+            delay_qvalue=True,
+            action_spec=env.full_action_spec_unbatched,
+        )
+        loss_module.set_keys(
+            state_action_value=("agents", "state_action_value"),
+            action=env.action_key,
+            reward=env.reward_key,
+            done=("agents", "done"),
+            terminated=("agents", "terminated"),
+        )
+    else:
+        loss_module = DiscreteSACLoss(
+            actor_network=policy,
+            qvalue_network=value_module,
+            delay_qvalue=True,
+            num_actions=env.full_action_spec_unbatched["agents", "action"].space.n,
+            action_space=env.full_action_spec_unbatched,
+        )
+        loss_module.set_keys(
+            action_value=("agents", "action_value"),
+            action=env.action_key,
+            reward=env.reward_key,
+            done=("agents", "done"),
+            terminated=("agents", "terminated"),
+        )
+
+    loss_module.make_value_estimator(ValueEstimators.TD0, gamma=cfg.loss.gamma)
+    target_net_updater = SoftUpdate(loss_module, eps=1 - cfg.loss.tau)
+
+    optim = torch.optim.Adam(loss_module.parameters(), cfg.train.lr)
+
+    # Logging
+    if cfg.logger.backend:
+        model_name = (
+            ("Het" if not cfg.model.shared_parameters else "")
+            + ("MA" if cfg.model.centralised_critic else "I")
+            + "SAC"
+        )
+        logger = init_logging(cfg, model_name)
+
+    total_time = 0
+    total_frames = 0
+    sampling_start = time.time()
+    for i, tensordict_data in enumerate(collector):
+        torchrl_logger.info(f"\nIteration {i}")
+
+        sampling_time = time.time() - sampling_start
+
+        current_frames = tensordict_data.numel()
+        total_frames += current_frames
+        data_view = tensordict_data.reshape(-1)
+        replay_buffer.extend(data_view)
+
+        training_tds = []
+        training_start = time.time()
+        for _ in range(cfg.train.num_epochs):
+            for _ in range(cfg.collector.frames_per_batch // cfg.train.minibatch_size):
+                subdata = replay_buffer.sample()
+                loss_vals = loss_module(subdata)
+                training_tds.append(loss_vals.detach())
+
+                loss_value = (
+                    loss_vals["loss_actor"]
+                    + loss_vals["loss_alpha"]
+                    + loss_vals["loss_qvalue"]
+                )
+
+                loss_value.backward()
+
+                total_norm = torch.nn.utils.clip_grad_norm_(
+                    loss_module.parameters(), cfg.train.max_grad_norm
+                )
+                training_tds[-1].set("grad_norm", total_norm.mean())
+
+                optim.step()
+                optim.zero_grad()
+                target_net_updater.step()
+
+        collector.update_policy_weights_()
+
+        training_time = time.time() - training_start
+
+        iteration_time = sampling_time + training_time
+        total_time += iteration_time
+        training_tds = torch.stack(training_tds)
+
+        # More logs
+        if cfg.logger.backend:
+            log_training(
+                logger,
+                training_tds,
+                tensordict_data,
+                sampling_time,
+                training_time,
+                total_time,
+                i,
+                current_frames,
+                total_frames,
+                step=i,
+            )
+
+        if (
+            cfg.eval.evaluation_episodes > 0
+            and i % cfg.eval.evaluation_interval == 0
+            and cfg.logger.backend
+        ):
+            evaluation_start = time.time()
+            with torch.no_grad(), set_exploration_type(ExplorationType.DETERMINISTIC):
+                env_test.frames = []
+                rollouts = env_test.rollout(
+                    max_steps=cfg.env.max_steps,
+                    policy=policy,
+                    callback=rendering_callback,
+                    auto_cast_to_device=True,
+                    break_when_any_done=False,
+                    # We are running vectorized evaluation we do not want it to stop when just one env is done
+                )
+
+                evaluation_time = time.time() - evaluation_start
+
+                log_evaluation(logger, rollouts, env_test, evaluation_time, step=i)
+
+        if cfg.logger.backend == "wandb":
+            logger.experiment.log({}, commit=True)
+        sampling_start = time.time()
+    collector.shutdown()
+    if not env.is_closed:
+        env.close()
+    if not env_test.is_closed:
+        env_test.close()
+
+
+if __name__ == "__main__":
+    train()

sota-implementations/multiagent/utils/__init__.py

@@ -0,0 +1,4 @@
+# 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.

sota-implementations/multiagent/utils/logging.py

@@ -0,0 +1,151 @@
+# 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 os
+
+import numpy as np
+import torch
+from tensordict import TensorDictBase
+from torchrl.envs.libs.vmas import VmasEnv
+from torchrl.record.loggers import generate_exp_name, get_logger, Logger
+from torchrl.record.loggers.wandb import WandbLogger
+
+
+def init_logging(cfg, model_name: str):
+    logger = get_logger(
+        logger_type=cfg.logger.backend,
+        logger_name=os.getcwd(),
+        experiment_name=generate_exp_name(cfg.env.scenario_name, model_name),
+        wandb_kwargs={
+            "group": cfg.logger.group_name or model_name,
+            "project": cfg.logger.project_name
+            or f"torchrl_example_{cfg.env.scenario_name}",
+        },
+    )
+    logger.log_hparams(cfg)
+    return logger
+
+
+def log_training(
+    logger: Logger,
+    training_td: TensorDictBase,
+    sampling_td: TensorDictBase,
+    sampling_time: float,
+    training_time: float,
+    total_time: float,
+    iteration: int,
+    current_frames: int,
+    total_frames: int,
+    step: int,
+):
+    if ("next", "agents", "reward") not in sampling_td.keys(True, True):
+        sampling_td.set(
+            ("next", "agents", "reward"),
+            sampling_td.get(("next", "reward"))
+            .expand(sampling_td.get("agents").shape)
+            .unsqueeze(-1),
+        )
+    if ("next", "agents", "episode_reward") not in sampling_td.keys(True, True):
+        sampling_td.set(
+            ("next", "agents", "episode_reward"),
+            sampling_td.get(("next", "episode_reward"))
+            .expand(sampling_td.get("agents").shape)
+            .unsqueeze(-1),
+        )
+
+    metrics_to_log = {
+        f"train/learner/{key}": value.mean().item()
+        for key, value in training_td.items()
+    }
+
+    if "info" in sampling_td.get("agents").keys():
+        metrics_to_log.update(
+            {
+                f"train/info/{key}": value.mean().item()
+                for key, value in sampling_td.get(("agents", "info")).items()
+            }
+        )
+
+    reward = sampling_td.get(("next", "agents", "reward")).mean(-2)  # Mean over agents
+    done = sampling_td.get(("next", "done"))
+    if done.ndim > reward.ndim:
+        done = done[..., 0, :]  # Remove expanded agent dim
+    episode_reward = sampling_td.get(("next", "agents", "episode_reward")).mean(-2)[
+        done
+    ]
+    metrics_to_log.update(
+        {
+            "train/reward/reward_min": reward.min().item(),
+            "train/reward/reward_mean": reward.mean().item(),
+            "train/reward/reward_max": reward.max().item(),
+            "train/reward/episode_reward_min": episode_reward.min().item(),
+            "train/reward/episode_reward_mean": episode_reward.mean().item(),
+            "train/reward/episode_reward_max": episode_reward.max().item(),
+            "train/sampling_time": sampling_time,
+            "train/training_time": training_time,
+            "train/iteration_time": training_time + sampling_time,
+            "train/total_time": total_time,
+            "train/training_iteration": iteration,
+            "train/current_frames": current_frames,
+            "train/total_frames": total_frames,
+        }
+    )
+    if isinstance(logger, WandbLogger):
+        logger.experiment.log(metrics_to_log, commit=False)
+    else:
+        for key, value in metrics_to_log.items():
+            logger.log_scalar(key.replace("/", "_"), value, step=step)
+
+    return metrics_to_log
+
+
+def log_evaluation(
+    logger: WandbLogger,
+    rollouts: TensorDictBase,
+    env_test: VmasEnv,
+    evaluation_time: float,
+    step: int,
+):
+    rollouts = list(rollouts.unbind(0))
+    for k, r in enumerate(rollouts):
+        next_done = r.get(("next", "done")).sum(
+            tuple(range(r.batch_dims, r.get(("next", "done")).ndim)),
+            dtype=torch.bool,
+        )
+        done_index = next_done.nonzero(as_tuple=True)[0][
+            0
+        ]  # First done index for this traj
+        rollouts[k] = r[: done_index + 1]
+
+    rewards = [td.get(("next", "agents", "reward")).sum(0).mean() for td in rollouts]
+    metrics_to_log = {
+        "eval/episode_reward_min": min(rewards),
+        "eval/episode_reward_max": max(rewards),
+        "eval/episode_reward_mean": sum(rewards) / len(rollouts),
+        "eval/episode_len_mean": sum([td.batch_size[0] for td in rollouts])
+        / len(rollouts),
+        "eval/evaluation_time": evaluation_time,
+    }
+
+    vid = torch.tensor(
+        np.transpose(env_test.frames[: rollouts[0].batch_size[0]], (0, 3, 1, 2)),
+        dtype=torch.uint8,
+    ).unsqueeze(0)
+
+    if isinstance(logger, WandbLogger):
+        import wandb
+
+        logger.experiment.log(metrics_to_log, commit=False)
+        logger.experiment.log(
+            {
+                "eval/video": wandb.Video(vid, fps=1 / env_test.world.dt, format="mp4"),
+            },
+            commit=False,
+        )
+    else:
+        for key, value in metrics_to_log.items():
+            logger.log_scalar(key.replace("/", "_"), value, step=step)
+        logger.log_video("eval_video", vid, step=step)

sota-implementations/multiagent/utils/utils.py

@@ -0,0 +1,43 @@
+# 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
+
+from tensordict import unravel_key
+from torchrl.envs import Transform
+
+
+def swap_last(source, dest):
+    source = unravel_key(source)
+    dest = unravel_key(dest)
+    if isinstance(source, str):
+        if isinstance(dest, str):
+            return dest
+        return dest[-1]
+    if isinstance(dest, str):
+        return source[:-1] + (dest,)
+    return source[:-1] + (dest[-1],)
+
+
+class DoneTransform(Transform):
+    """Expands the 'done' entries (incl. terminated) to match the reward shape.
+
+    Can be appended to a replay buffer or a collector.
+    """
+
+    def __init__(self, reward_key, done_keys):
+        super().__init__()
+        self.reward_key = reward_key
+        self.done_keys = done_keys
+
+    def forward(self, tensordict):
+        for done_key in self.done_keys:
+            new_name = swap_last(self.reward_key, done_key)
+            tensordict.set(
+                ("next", new_name),
+                tensordict.get(("next", done_key))
+                .unsqueeze(-1)
+                .expand(tensordict.get(("next", self.reward_key)).shape),
+            )
+        return tensordict

sota-implementations/ppo/README.md

@@ -0,0 +1,29 @@
+## Reproducing Proximal Policy Optimization (PPO) Algorithm Results
+
+This repository contains scripts that enable training agents using the Proximal Policy Optimization (PPO) Algorithm on MuJoCo and Atari environments. We follow the original paper [Proximal Policy Optimization Algorithms](https://arxiv.org/abs/1707.06347) by Schulman et al. (2017) to implement the PPO algorithm but introduce the improvement of computing the Generalised Advantage Estimator (GAE) at every epoch.
+
+
+## Examples Structure
+
+Please note that each example is independent of each other for the sake of simplicity. Each example contains the following files:
+
+1. **Main Script:** The definition of algorithm components and the training loop can be found in the main script  (e.g. ppo_atari.py).
+
+2. **Utils File:** A utility file is provided to contain various helper functions, generally to create the environment and the models (e.g. utils_atari.py).
+
+3. **Configuration File:** This file includes default hyperparameters specified in the original paper. Users can modify these hyperparameters to customize their experiments  (e.g. config_atari.yaml).
+
+
+## Running the Examples
+
+You can execute the PPO algorithm on Atari environments by running the following command:
+
+```bash
+python ppo_atari.py
+```
+
+You can execute the PPO algorithm on MuJoCo environments by running the following command:
+
+```bash
+python ppo_mujoco.py
+```