torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

sota-implementations/gail/gail.py

@@ -0,0 +1,327 @@
+# 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.
+"""GAIL Example.
+
+This is a self-contained example of an offline GAIL training script.
+
+The helper functions for gail are coded in the gail_utils.py and helper functions for ppo in ppo_utils.
+
+"""
+from __future__ import annotations
+
+import warnings
+
+import hydra
+import numpy as np
+import torch
+import tqdm
+from gail_utils import log_metrics, make_gail_discriminator, make_offline_replay_buffer
+from ppo_utils import eval_model, make_env, make_ppo_models
+from tensordict.nn import CudaGraphModule
+from torchrl._utils import compile_with_warmup, get_available_device, timeit
+from torchrl.collectors import SyncDataCollector
+from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer
+from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
+from torchrl.envs import set_gym_backend
+from torchrl.envs.utils import ExplorationType, set_exploration_type
+from torchrl.objectives import ClipPPOLoss, GAILLoss, group_optimizers
+from torchrl.objectives.value.advantages import GAE
+from torchrl.record import VideoRecorder
+from torchrl.record.loggers import generate_exp_name, get_logger
+
+torch.set_float32_matmul_precision("high")
+
+
+@hydra.main(config_path="", config_name="config")
+def main(cfg: DictConfig):  # noqa: F821
+    set_gym_backend(cfg.env.backend).set()
+
+    device = (
+        torch.device(cfg.gail.device) if cfg.gail.device else get_available_device()
+    )
+    num_mini_batches = (
+        cfg.ppo.collector.frames_per_batch // cfg.ppo.loss.mini_batch_size
+    )
+    total_network_updates = (
+        (cfg.ppo.collector.total_frames // cfg.ppo.collector.frames_per_batch)
+        * cfg.ppo.loss.ppo_epochs
+        * num_mini_batches
+    )
+
+    # Create logger
+    exp_name = generate_exp_name("Gail", cfg.logger.exp_name)
+    logger = None
+    if cfg.logger.backend:
+        logger = get_logger(
+            logger_type=cfg.logger.backend,
+            logger_name="gail_logging",
+            experiment_name=exp_name,
+            wandb_kwargs={
+                "mode": cfg.logger.mode,
+                "config": dict(cfg),
+                "project": cfg.logger.project_name,
+                "group": cfg.logger.group_name,
+            },
+        )
+
+    # Set seeds
+    torch.manual_seed(cfg.env.seed)
+    np.random.seed(cfg.env.seed)
+
+    # Create models (check utils_mujoco.py)
+    actor, critic = make_ppo_models(
+        cfg.env.env_name, compile=cfg.compile.compile, device=device
+    )
+
+    # Create data buffer
+    data_buffer = TensorDictReplayBuffer(
+        storage=LazyTensorStorage(
+            cfg.ppo.collector.frames_per_batch,
+            device=device,
+            compilable=cfg.compile.compile,
+        ),
+        sampler=SamplerWithoutReplacement(),
+        batch_size=cfg.ppo.loss.mini_batch_size,
+        compilable=cfg.compile.compile,
+    )
+
+    # Create loss and adv modules
+    adv_module = GAE(
+        gamma=cfg.ppo.loss.gamma,
+        lmbda=cfg.ppo.loss.gae_lambda,
+        value_network=critic,
+        average_gae=False,
+        device=device,
+    )
+
+    loss_module = ClipPPOLoss(
+        actor_network=actor,
+        critic_network=critic,
+        clip_epsilon=cfg.ppo.loss.clip_epsilon,
+        loss_critic_type=cfg.ppo.loss.loss_critic_type,
+        entropy_coeff=cfg.ppo.loss.entropy_coeff,
+        critic_coeff=cfg.ppo.loss.critic_coeff,
+        normalize_advantage=True,
+    )
+
+    # Create optimizers
+    actor_optim = torch.optim.Adam(
+        actor.parameters(), lr=torch.tensor(cfg.ppo.optim.lr, device=device), eps=1e-5
+    )
+    critic_optim = torch.optim.Adam(
+        critic.parameters(), lr=torch.tensor(cfg.ppo.optim.lr, device=device), eps=1e-5
+    )
+    optim = group_optimizers(actor_optim, critic_optim)
+    del actor_optim, critic_optim
+
+    compile_mode = None
+    if cfg.compile.compile:
+        compile_mode = cfg.compile.compile_mode
+        if compile_mode in ("", None):
+            if cfg.compile.cudagraphs:
+                compile_mode = "default"
+            else:
+                compile_mode = "reduce-overhead"
+
+    # Create collector
+    collector = SyncDataCollector(
+        create_env_fn=make_env(cfg.env.env_name, device),
+        policy=actor,
+        frames_per_batch=cfg.ppo.collector.frames_per_batch,
+        total_frames=cfg.ppo.collector.total_frames,
+        device=device,
+        max_frames_per_traj=-1,
+        compile_policy={"mode": compile_mode} if compile_mode is not None else False,
+        cudagraph_policy={"warmup": 10} if cfg.compile.cudagraphs else False,
+    )
+
+    # Create replay buffer
+    replay_buffer = make_offline_replay_buffer(cfg.replay_buffer)
+
+    # Create Discriminator
+    discriminator = make_gail_discriminator(cfg, collector.env, device)
+
+    # Create loss
+    discriminator_loss = GAILLoss(
+        discriminator,
+        use_grad_penalty=cfg.gail.use_grad_penalty,
+        gp_lambda=cfg.gail.gp_lambda,
+    )
+
+    # Create optimizer
+    discriminator_optim = torch.optim.Adam(
+        params=discriminator.parameters(), lr=cfg.gail.lr
+    )
+
+    # Create test environment
+    logger_video = cfg.logger.video
+    test_env = make_env(cfg.env.env_name, device, from_pixels=logger_video)
+    if logger_video:
+        test_env = test_env.append_transform(
+            VideoRecorder(logger, tag="rendering/test", in_keys=["pixels"])
+        )
+    test_env.eval()
+    num_network_updates = torch.zeros((), dtype=torch.int64, device=device)
+
+    def update(data, expert_data, num_network_updates=num_network_updates):
+        # Add collector data to expert data
+        expert_data.set(
+            discriminator_loss.tensor_keys.collector_action,
+            data["action"][: expert_data.batch_size[0]],
+        )
+        expert_data.set(
+            discriminator_loss.tensor_keys.collector_observation,
+            data["observation"][: expert_data.batch_size[0]],
+        )
+        d_loss = discriminator_loss(expert_data)
+
+        # Backward pass
+        d_loss.get("loss").backward()
+        discriminator_optim.step()
+        discriminator_optim.zero_grad(set_to_none=True)
+
+        # Compute discriminator reward
+        with torch.no_grad():
+            data = discriminator(data)
+        d_rewards = -torch.log(1 - data["d_logits"] + 1e-8)
+
+        # Set discriminator rewards to tensordict
+        data.set(("next", "reward"), d_rewards)
+
+        # Update PPO
+        for _ in range(cfg_loss_ppo_epochs):
+            # Compute GAE
+            with torch.no_grad():
+                data = adv_module(data)
+            data_reshape = data.reshape(-1)
+
+            # Update the data buffer
+            data_buffer.empty()
+            data_buffer.extend(data_reshape)
+
+            for batch in data_buffer:
+                optim.zero_grad(set_to_none=True)
+
+                # Linearly decrease the learning rate and clip epsilon
+                alpha = torch.ones((), device=device)
+                if cfg_optim_anneal_lr:
+                    alpha = 1 - (num_network_updates / total_network_updates)
+                    for group in optim.param_groups:
+                        group["lr"] = cfg_optim_lr * alpha
+                if cfg_loss_anneal_clip_eps:
+                    loss_module.clip_epsilon.copy_(cfg_loss_clip_epsilon * alpha)
+                num_network_updates += 1
+
+                # Forward pass PPO loss
+                loss = loss_module(batch)
+                critic_loss = loss["loss_critic"]
+                actor_loss = loss["loss_objective"] + loss["loss_entropy"]
+
+                # Backward pass
+                (actor_loss + critic_loss).backward()
+
+                # Update the networks
+                optim.step()
+        return {"dloss": d_loss, "alpha": alpha}
+
+    if cfg.compile.compile:
+        update = compile_with_warmup(update, warmup=2, mode=compile_mode)
+    if cfg.compile.cudagraphs:
+        warnings.warn(
+            "CudaGraphModule is experimental and may lead to silently wrong results. Use with caution.",
+            category=UserWarning,
+        )
+        update = CudaGraphModule(update, warmup=50)
+
+    # Training loop
+    collected_frames = 0
+    pbar = tqdm.tqdm(total=cfg.ppo.collector.total_frames)
+
+    # extract cfg variables
+    cfg_loss_ppo_epochs = cfg.ppo.loss.ppo_epochs
+    cfg_optim_anneal_lr = cfg.ppo.optim.anneal_lr
+    cfg_optim_lr = cfg.ppo.optim.lr
+    cfg_loss_anneal_clip_eps = cfg.ppo.loss.anneal_clip_epsilon
+    cfg_loss_clip_epsilon = cfg.ppo.loss.clip_epsilon
+    cfg_logger_test_interval = cfg.logger.test_interval
+    cfg_logger_num_test_episodes = cfg.logger.num_test_episodes
+
+    total_iter = len(collector)
+    collector_iter = iter(collector)
+    for i in range(total_iter):
+
+        timeit.printevery(1000, total_iter, erase=True)
+
+        with timeit("collection"):
+            data = next(collector_iter)
+
+        metrics_to_log = {}
+        frames_in_batch = data.numel()
+        collected_frames += frames_in_batch
+        pbar.update(data.numel())
+
+        with timeit("rb - sample expert"):
+            # Get expert data
+            expert_data = replay_buffer.sample()
+            expert_data = expert_data.to(device)
+
+        with timeit("update"):
+            torch.compiler.cudagraph_mark_step_begin()
+            metadata = update(data, expert_data)
+        d_loss = metadata["dloss"]
+        alpha = metadata["alpha"]
+
+        # Get training rewards and episode lengths
+        episode_rewards = data["next", "episode_reward"][data["next", "done"]]
+        if len(episode_rewards) > 0:
+            episode_length = data["next", "step_count"][data["next", "done"]]
+
+            metrics_to_log.update(
+                {
+                    "train/reward": episode_rewards.mean().item(),
+                    "train/episode_length": episode_length.sum().item()
+                    / len(episode_length),
+                }
+            )
+
+        metrics_to_log.update(
+            {
+                "train/discriminator_loss": d_loss["loss"],
+                "train/lr": alpha * cfg_optim_lr,
+                "train/clip_epsilon": (
+                    alpha * cfg_loss_clip_epsilon
+                    if cfg_loss_anneal_clip_eps
+                    else cfg_loss_clip_epsilon
+                ),
+            }
+        )
+
+        # evaluation
+        with torch.no_grad(), set_exploration_type(
+            ExplorationType.DETERMINISTIC
+        ), timeit("eval"):
+            if ((i - 1) * frames_in_batch) // cfg_logger_test_interval < (
+                i * frames_in_batch
+            ) // cfg_logger_test_interval:
+                actor.eval()
+                test_rewards = eval_model(
+                    actor, test_env, num_episodes=cfg_logger_num_test_episodes
+                )
+                metrics_to_log.update(
+                    {
+                        "eval/reward": test_rewards.mean(),
+                    }
+                )
+                actor.train()
+        if logger is not None:
+            metrics_to_log.update(timeit.todict(prefix="time"))
+            metrics_to_log["time/speed"] = pbar.format_dict["rate"]
+            log_metrics(logger, metrics_to_log, i)
+
+    pbar.close()
+
+
+if __name__ == "__main__":
+    main()

sota-implementations/gail/gail_utils.py

@@ -0,0 +1,68 @@
+# 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 torch.nn as nn
+import torch.optim
+from torchrl.data.datasets.d4rl import D4RLExperienceReplay
+from torchrl.data.replay_buffers import SamplerWithoutReplacement
+from torchrl.envs import DoubleToFloat
+from torchrl.modules import SafeModule
+
+
+# ====================================================================
+# Offline 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,
+        sampler=SamplerWithoutReplacement(drop_last=False),
+        prefetch=4,
+        direct_download=True,
+    )
+
+    data.append_transform(DoubleToFloat())
+
+    return data
+
+
+def make_gail_discriminator(cfg, train_env, device="cpu"):
+    """Make GAIL discriminator."""
+
+    state_dim = train_env.observation_spec["observation"].shape[0]
+    action_dim = train_env.action_spec.shape[0]
+
+    hidden_dim = cfg.gail.hidden_dim
+
+    # Define Discriminator Network
+    class Discriminator(nn.Module):
+        def __init__(self, state_dim, action_dim):
+            super().__init__()
+            self.fc1 = nn.Linear(state_dim + action_dim, hidden_dim)
+            self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+            self.fc3 = nn.Linear(hidden_dim, 1)
+
+        def forward(self, state, action):
+            x = torch.cat([state, action], dim=1)
+            x = torch.relu(self.fc1(x))
+            x = torch.relu(self.fc2(x))
+            return torch.sigmoid(self.fc3(x))
+
+    d_module = SafeModule(
+        module=Discriminator(state_dim, action_dim),
+        in_keys=["observation", "action"],
+        out_keys=["d_logits"],
+    )
+    return d_module.to(device)
+
+
+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)

sota-implementations/gail/ppo_utils.py

@@ -0,0 +1,157 @@
+# 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 torch.nn
+import torch.optim
+
+from tensordict.nn import AddStateIndependentNormalScale, TensorDictModule
+from torchrl.envs import (
+    ClipTransform,
+    DoubleToFloat,
+    ExplorationType,
+    RewardSum,
+    StepCounter,
+    TransformedEnv,
+    VecNorm,
+)
+from torchrl.envs.libs.gym import GymEnv
+from torchrl.modules import MLP, ProbabilisticActor, TanhNormal, ValueOperator
+from torchrl.record import VideoRecorder
+
+
+# ====================================================================
+# Environment utils
+# --------------------------------------------------------------------
+
+
+def make_env(env_name="HalfCheetah-v4", device="cpu", from_pixels: bool = False):
+    env = GymEnv(env_name, device=device, from_pixels=from_pixels, pixels_only=False)
+    env = TransformedEnv(env)
+    env.append_transform(VecNorm(in_keys=["observation"], decay=0.99999, eps=1e-2))
+    env.append_transform(ClipTransform(in_keys=["observation"], low=-10, high=10))
+    env.append_transform(RewardSum())
+    env.append_transform(StepCounter())
+    env.append_transform(DoubleToFloat(in_keys=["observation"]))
+    return env
+
+
+# ====================================================================
+# Model utils
+# --------------------------------------------------------------------
+
+
+def make_ppo_models_state(proof_environment, compile, device):
+
+    # Define input shape
+    input_shape = proof_environment.observation_spec["observation"].shape
+
+    # Define policy output distribution class
+    num_outputs = proof_environment.action_spec_unbatched.shape[-1]
+    distribution_class = TanhNormal
+    distribution_kwargs = {
+        "low": proof_environment.action_spec_unbatched.space.low.to(device),
+        "high": proof_environment.action_spec_unbatched.space.high.to(device),
+        "tanh_loc": False,
+        # "safe_tanh": not compile,
+    }
+
+    # Define policy architecture
+    policy_mlp = MLP(
+        in_features=input_shape[-1],
+        activation_class=torch.nn.Tanh,
+        out_features=num_outputs,  # predict only loc
+        num_cells=[64, 64],
+        device=device,
+    )
+
+    # Initialize policy weights
+    for layer in policy_mlp.modules():
+        if isinstance(layer, torch.nn.Linear):
+            torch.nn.init.orthogonal_(layer.weight, 1.0)
+            layer.bias.data.zero_()
+
+    # Add state-independent normal scale
+    policy_mlp = torch.nn.Sequential(
+        policy_mlp,
+        AddStateIndependentNormalScale(
+            proof_environment.action_spec_unbatched.shape[-1],
+            scale_lb=1e-8,
+            device=device,
+        ),
+    )
+
+    # Add probabilistic sampling of the actions
+    policy_module = ProbabilisticActor(
+        TensorDictModule(
+            module=policy_mlp,
+            in_keys=["observation"],
+            out_keys=["loc", "scale"],
+        ),
+        in_keys=["loc", "scale"],
+        spec=proof_environment.full_action_spec_unbatched.to(device),
+        distribution_class=distribution_class,
+        distribution_kwargs=distribution_kwargs,
+        return_log_prob=True,
+        default_interaction_type=ExplorationType.RANDOM,
+    )
+
+    # Define value architecture
+    value_mlp = MLP(
+        in_features=input_shape[-1],
+        activation_class=torch.nn.Tanh,
+        out_features=1,
+        num_cells=[64, 64],
+        device=device,
+    )
+
+    # Initialize value weights
+    for layer in value_mlp.modules():
+        if isinstance(layer, torch.nn.Linear):
+            torch.nn.init.orthogonal_(layer.weight, 0.01)
+            layer.bias.data.zero_()
+
+    # Define value module
+    value_module = ValueOperator(
+        value_mlp,
+        in_keys=["observation"],
+    )
+
+    return policy_module, value_module
+
+
+def make_ppo_models(env_name, compile, device):
+    proof_environment = make_env(env_name, device=device)
+    actor, critic = make_ppo_models_state(
+        proof_environment, compile=compile, device=device
+    )
+    return actor, critic
+
+
+# ====================================================================
+# Evaluation utils
+# --------------------------------------------------------------------
+
+
+def dump_video(module):
+    if isinstance(module, VideoRecorder):
+        module.dump()
+
+
+def eval_model(actor, test_env, num_episodes=3):
+    test_rewards = []
+    for _ in range(num_episodes):
+        td_test = test_env.rollout(
+            policy=actor,
+            auto_reset=True,
+            auto_cast_to_device=True,
+            break_when_any_done=True,
+            max_steps=10_000_000,
+        )
+        reward = td_test["next", "episode_reward"][td_test["next", "done"]]
+        test_rewards.append(reward.cpu())
+        test_env.apply(dump_video)
+    del td_test
+    return torch.cat(test_rewards, 0).mean()