torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

sota-implementations/decision_transformer/lamb.py

@@ -0,0 +1,167 @@
+# 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.
+# Lamb optimizer directly copied from https://github.com/facebookresearch/online-dt
+from __future__ import annotations
+
+import math
+
+import torch
+from torch.optim import Optimizer
+
+
+class Lamb(Optimizer):
+    """Implements a pure pytorch variant of FuseLAMB (NvLamb variant) optimizer from apex.optimizers.FusedLAMB
+    reference: https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/Transformer-XL/pytorch/lamb.py
+    LAMB was proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining parameter groups.
+        lr (:obj:`float`, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its norm. (default: (0.9, 0.999))
+        eps (:obj:`float`, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        weight_decay (:obj:`float`, optional): weight decay (L2 penalty) (default: 0)
+        grad_averaging (bool, optional): whether apply (1-beta2) to grad when
+            calculating running averages of gradient. (default: True)
+        max_grad_norm (:obj:`float`, optional): value used to clip global grad norm (default: 1.0)
+        trust_clip (bool): enable LAMBC trust ratio clipping (default: False)
+        always_adapt (boolean, optional): Apply adaptive learning rate to 0.0
+            weight decay parameter (default: False)
+    .. _Large Batch Optimization for Deep Learning - Training BERT in 76 minutes:
+        https://arxiv.org/abs/1904.00962
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        bias_correction=True,
+        betas=(0.9, 0.999),
+        eps=1e-6,
+        weight_decay=0.01,
+        grad_averaging=True,
+        max_grad_norm=1.0,
+        trust_clip=False,
+        always_adapt=False,
+    ):
+        defaults = {
+            "lr": lr,
+            "bias_correction": bias_correction,
+            "betas": betas,
+            "eps": eps,
+            "weight_decay": weight_decay,
+            "grad_averaging": grad_averaging,
+            "max_grad_norm": max_grad_norm,
+            "trust_clip": trust_clip,
+            "always_adapt": always_adapt,
+        }
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        device = self.param_groups[0]["params"][0].device
+        one_tensor = torch.tensor(
+            1.0, device=device
+        )  # because torch.where doesn't handle scalars correctly
+        global_grad_norm = torch.zeros(1, device=device)
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError(
+                        "Lamb does not support sparse gradients, consider SparseAdam instead."
+                    )
+                global_grad_norm.add_(grad.pow(2).sum())
+
+        global_grad_norm = torch.sqrt(global_grad_norm)
+        # FIXME it'd be nice to remove explicit tensor conversion of scalars when torch.where promotes
+        # scalar types properly https://github.com/pytorch/pytorch/issues/9190
+        max_grad_norm = torch.tensor(self.defaults["max_grad_norm"], device=device)
+        clip_global_grad_norm = torch.where(
+            global_grad_norm > max_grad_norm,
+            global_grad_norm / max_grad_norm,
+            one_tensor,
+        )
+
+        for group in self.param_groups:
+            bias_correction = 1 if group["bias_correction"] else 0
+            beta1, beta2 = group["betas"]
+            grad_averaging = 1 if group["grad_averaging"] else 0
+            beta3 = 1 - beta1 if grad_averaging else 1.0
+
+            # assume same step across group now to simplify things
+            # per parameter step can be easily support by making it tensor, or pass list into kernel
+            if "step" in group:
+                group["step"] += 1
+            else:
+                group["step"] = 1
+
+            if bias_correction:
+                bias_correction1 = 1 - beta1 ** group["step"]
+                bias_correction2 = 1 - beta2 ** group["step"]
+            else:
+                bias_correction1, bias_correction2 = 1.0, 1.0
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.div_(clip_global_grad_norm)
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    # Exponential moving average of gradient valuesa
+                    state["exp_avg"] = torch.zeros_like(p)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(p)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(grad, alpha=beta3)  # m_t
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)  # v_t
+
+                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(
+                    group["eps"]
+                )
+                update = (exp_avg / bias_correction1).div_(denom)
+
+                weight_decay = group["weight_decay"]
+                if weight_decay != 0:
+                    update.add_(p, alpha=weight_decay)
+
+                if weight_decay != 0 or group["always_adapt"]:
+                    # Layer-wise LR adaptation. By default, skip adaptation on parameters that are
+                    # excluded from weight decay, unless always_adapt == True, then always enabled.
+                    w_norm = p.norm(2.0)
+                    g_norm = update.norm(2.0)
+                    # FIXME nested where required since logical and/or not working in PT XLA
+                    trust_ratio = torch.where(
+                        w_norm > 0,
+                        torch.where(g_norm > 0, w_norm / g_norm, one_tensor),
+                        one_tensor,
+                    )
+                    if group["trust_clip"]:
+                        # LAMBC trust clipping, upper bound fixed at one
+                        trust_ratio = torch.minimum(trust_ratio, one_tensor)
+                    update.mul_(trust_ratio)
+
+                p.add_(update, alpha=-group["lr"])
+
+        return loss

sota-implementations/decision_transformer/online_dt.py

@@ -0,0 +1,178 @@
+# 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.
+"""Online Decision Transformer Example.
+This is a self-contained example of an Online Decision Transformer training script.
+The helper functions are coded in the utils.py associated with this script.
+"""
+from __future__ import annotations
+
+import warnings
+
+import hydra
+import numpy as np
+import torch
+import tqdm
+from tensordict.nn import CudaGraphModule
+from torchrl._utils import get_available_device, logger as torchrl_logger, timeit
+from torchrl.envs.libs.gym import set_gym_backend
+from torchrl.envs.utils import ExplorationType, set_exploration_type
+from torchrl.modules.tensordict_module import DecisionTransformerInferenceWrapper
+from torchrl.record import VideoRecorder
+from utils import (
+    dump_video,
+    log_metrics,
+    make_env,
+    make_logger,
+    make_odt_loss,
+    make_odt_model,
+    make_odt_optimizer,
+    make_offline_replay_buffer,
+)
+
+
+@hydra.main(config_path="", config_name="odt_config", version_base="1.1")
+def main(cfg: DictConfig):  # noqa: F821
+    set_gym_backend(cfg.env.backend).set()
+
+    model_device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
+
+    # Set seeds
+    torch.manual_seed(cfg.env.seed)
+    np.random.seed(cfg.env.seed)
+
+    # Create logger
+    logger = make_logger(cfg)
+
+    # Create offline replay buffer
+    offline_buffer, obs_loc, obs_std = make_offline_replay_buffer(
+        cfg.replay_buffer, cfg.env.reward_scaling
+    )
+
+    # Create test environment
+    test_env = make_env(cfg.env, obs_loc, obs_std, from_pixels=cfg.logger.video)
+    if cfg.logger.video:
+        test_env = test_env.append_transform(
+            VideoRecorder(logger, tag="rendered", in_keys=["pixels"])
+        )
+
+    # Create policy model
+    policy = make_odt_model(cfg, device=model_device)
+
+    # Create loss
+    loss_module = make_odt_loss(cfg.loss, policy)
+
+    # Create optimizer
+    transformer_optim, temperature_optim, scheduler = make_odt_optimizer(
+        cfg.optim, loss_module
+    )
+
+    # Create inference policy
+    inference_policy = DecisionTransformerInferenceWrapper(
+        policy=policy,
+        inference_context=cfg.env.inference_context,
+        device=model_device,
+    )
+    inference_policy.set_tensor_keys(
+        observation="observation_cat",
+        action="action_cat",
+        return_to_go="return_to_go_cat",
+    )
+
+    def update(data):
+        transformer_optim.zero_grad(set_to_none=True)
+        temperature_optim.zero_grad(set_to_none=True)
+        # Compute loss
+        loss_vals = loss_module(data.to(model_device))
+        transformer_loss = loss_vals["loss_log_likelihood"] + loss_vals["loss_entropy"]
+        temperature_loss = loss_vals["loss_alpha"]
+
+        (temperature_loss + transformer_loss).backward()
+        torch.nn.utils.clip_grad_norm_(policy.parameters(), clip_grad)
+
+        transformer_optim.step()
+        temperature_optim.step()
+
+        return loss_vals.detach()
+
+    if cfg.compile.compile:
+        compile_mode = cfg.compile.compile_mode
+        if compile_mode in ("", None):
+            compile_mode = "default"
+        update = torch.compile(update, mode=compile_mode, dynamic=False)
+    if cfg.compile.cudagraphs:
+        warnings.warn(
+            "CudaGraphModule is experimental and may lead to silently wrong results. Use with caution.",
+            category=UserWarning,
+        )
+        if cfg.optim.optimizer == "lamb":
+            raise ValueError(
+                "cudagraphs isn't compatible with the Lamb optimizer. Use optim.optimizer=Adam instead."
+            )
+        update = CudaGraphModule(update, warmup=50)
+
+    pbar = tqdm.tqdm(total=cfg.optim.pretrain_gradient_steps)
+
+    pretrain_gradient_steps = cfg.optim.pretrain_gradient_steps
+    clip_grad = cfg.optim.clip_grad
+    eval_steps = cfg.logger.eval_steps
+    pretrain_log_interval = cfg.logger.pretrain_log_interval
+    reward_scaling = cfg.env.reward_scaling
+
+    torchrl_logger.info(" ***Pretraining*** ")
+    # Pretraining
+    for i in range(pretrain_gradient_steps):
+        timeit.printevery(1000, pretrain_gradient_steps, erase=True)
+        pbar.update(1)
+        with timeit("sample"):
+            # Sample data
+            data = offline_buffer.sample()
+
+        with timeit("update"):
+            torch.compiler.cudagraph_mark_step_begin()
+            loss_vals = update(data.to(model_device))
+
+        scheduler.step()
+
+        # Log metrics
+        metrics_to_log = {
+            "train/loss_log_likelihood": loss_vals["loss_log_likelihood"],
+            "train/loss_entropy": loss_vals["loss_entropy"],
+            "train/loss_alpha": loss_vals["loss_alpha"],
+            "train/alpha": loss_vals["alpha"],
+            "train/entropy": loss_vals["entropy"],
+        }
+
+        # Evaluation
+        with torch.no_grad(), set_exploration_type(
+            ExplorationType.DETERMINISTIC
+        ), timeit("eval"):
+            inference_policy.eval()
+            if i % pretrain_log_interval == 0:
+                eval_td = test_env.rollout(
+                    max_steps=eval_steps,
+                    policy=inference_policy,
+                    auto_cast_to_device=True,
+                    break_when_any_done=False,
+                )
+                test_env.apply(dump_video)
+                inference_policy.train()
+            metrics_to_log["eval/reward"] = (
+                eval_td["next", "reward"].sum(1).mean().item() / reward_scaling
+            )
+
+        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 not test_env.is_closed:
+        test_env.close()
+
+
+if __name__ == "__main__":
+    main()