egogym 0.1.0__py3-none-any.whl

baselines/rum/models/bet/bet.py

@@ -0,0 +1,347 @@
+import logging
+from enum import Enum
+from typing import Dict, Optional, Tuple
+
+import accelerate
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import tqdm
+
+from baselines.rum.models.bet.gpt import GPT
+from baselines.rum.models.bet.utils import MLP
+
+
+GENERATOR_SEED_FIXED = 123456789
+
+
+class KMeansDiscretizer:
+    """
+    Simplified and modified version of KMeans algorithm from sklearn.
+    We initialize this with a fixed seed to ensure that on each GPU we come up with the same
+    clusters.
+    """
+
+    def __init__(
+        self,
+        num_bins: int = 100,
+        kmeans_iters: int = 50,
+    ):
+        super().__init__()
+        self.n_bins = num_bins
+        self.kmeans_iters = kmeans_iters
+
+    def fit(self, input_actions: torch.Tensor) -> None:
+        self.bin_centers = KMeansDiscretizer._kmeans(
+            input_actions, ncluster=self.n_bins, niter=self.kmeans_iters
+        )
+
+    @classmethod
+    def _kmeans(cls, x: torch.Tensor, ncluster: int = 512, niter: int = 50):
+        """
+        Simple k-means clustering algorithm adapted from Karpathy's minGPT libary
+        https://github.com/karpathy/minGPT/blob/master/play_image.ipynb
+        """
+        N, D = x.size()
+        generator = torch.Generator()
+        generator.manual_seed(GENERATOR_SEED_FIXED)
+
+        c = x[
+            torch.randperm(N, generator=generator)[:ncluster]
+        ]  # init clusters at random, with a fixed seed
+
+        pbar = tqdm.trange(niter)
+        pbar.set_description("K-means clustering")
+        for i in pbar:
+            # assign all pixels to the closest codebook element
+            a = ((x[:, None, :] - c[None, :, :]) ** 2).sum(-1).argmin(1)
+            # move each codebook element to be the mean of the pixels that assigned to it
+            c = torch.stack([x[a == k].mean(0) for k in range(ncluster)])
+            # re-assign any poorly positioned codebook elements
+            nanix = torch.any(torch.isnan(c), dim=1)
+            ndead = nanix.sum().item()
+            if ndead:
+                tqdm.tqdm.write(
+                    "done step %d/%d, re-initialized %d dead clusters"
+                    % (i + 1, niter, ndead)
+                )
+            c[nanix] = x[
+                torch.randperm(N, generator=generator)[:ndead]
+            ]  # re-init dead clusters
+        return c
+
+
+class BehaviorTransformer(nn.Module):
+    GOAL_SPEC = Enum("GOAL_SPEC", "concat stack unconditional")
+
+    def __init__(
+        self,
+        obs_dim: int,
+        act_dim: int,
+        goal_dim: int,
+        gpt_model: GPT,
+        num_extra_predicted_actions: Optional[int] = None,
+        trainable_obs_padding: bool = True,
+        n_clusters: int = 32,
+        kmeans_fit_steps: int = 500,
+        kmeans_iters: int = 50,
+        offset_loss_multiplier: float = 1.0e3,
+        offset_distance_metric: str = "L2",
+        gamma: float = 2.0,
+        **kwargs,
+    ):
+        super().__init__()
+        self._obs_dim = obs_dim
+        self._act_dim = act_dim
+        self._goal_dim = goal_dim
+        self._num_extra_predicted_actions = num_extra_predicted_actions
+        # Gradient-free, all zeros if we don't want to train this.
+        self._obs_padding = nn.Parameter(
+            trainable_obs_padding * torch.randn(obs_dim),
+            requires_grad=trainable_obs_padding,
+        )
+
+        if goal_dim <= 0:
+            self._cbet_method = self.GOAL_SPEC.unconditional
+        elif obs_dim == goal_dim:
+            self._cbet_method = self.GOAL_SPEC.concat
+        else:
+            self._cbet_method = self.GOAL_SPEC.stack
+
+        self._gpt_model = gpt_model
+        # For now, we assume the number of clusters is given.
+        assert n_clusters > 0
+        self._K = n_clusters
+        self._kmeans_fit_steps = kmeans_fit_steps
+        self._clustering_algo = KMeansDiscretizer(
+            num_bins=n_clusters, kmeans_iters=kmeans_iters
+        )
+        self._current_steps = 0
+        self._map_to_cbet_preds = MLP(
+            in_channels=gpt_model.config.output_dim,
+            hidden_channels=[(act_dim + 1) * n_clusters],
+        )
+        self._collected_actions = []
+        self._have_fit_kmeans = False
+        self._offset_loss_multiplier = offset_loss_multiplier
+        # Placeholder for the cluster centers.
+        generator = torch.Generator()
+        generator.manual_seed(GENERATOR_SEED_FIXED)
+        self.register_buffer(
+            "_cluster_centers",
+            torch.randn(
+                (n_clusters, act_dim), generator=generator, dtype=torch.float32
+            ),
+        )
+        self._criterion = FocalLoss(gamma=gamma, reduction="none")
+        self._offset_criterion = (
+            nn.MSELoss(reduction="none")
+            if offset_distance_metric == "L2"
+            else nn.L1Loss(reduction="none")
+        )
+        self._accelerator = accelerate.Accelerator()
+
+    def _load_from_state_dict(self, *args, **kwargs):
+        # Don't fit kmeans if we are loading from a state dict.
+        self._current_steps = self._kmeans_fit_steps
+        self._have_fit_kmeans = True
+        return super()._load_from_state_dict(*args, **kwargs)
+
+    def forward(
+        self,
+        obs_seq: torch.Tensor,
+        goal_seq: Optional[torch.Tensor],
+        action_seq: Optional[torch.Tensor],
+        padding_seq: Optional[torch.Tensor],
+        predict_with_offset: bool = True,
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        if self._current_steps == 0:
+            self._cluster_centers = self._cluster_centers.to(obs_seq.device)
+        if self._current_steps < self._kmeans_fit_steps and action_seq is not None:
+            self._current_steps += 1
+            self._fit_kmeans(obs_seq, goal_seq, action_seq, padding_seq)
+        return self._predict(
+            obs_seq,
+            goal_seq,
+            action_seq,
+            padding_seq,
+            predict_with_offset=predict_with_offset,
+        )
+
+    def _fit_kmeans(
+        self,
+        obs_seq: torch.Tensor,
+        goal_seq: Optional[torch.Tensor],
+        action_seq: Optional[torch.Tensor],
+        padding_seq: Optional[torch.Tensor],
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        assert self._current_steps <= self._kmeans_fit_steps
+        if self._current_steps == 1:
+            self._cluster_centers = self._cluster_centers.to(action_seq.device)
+
+        all_action_seq = self._accelerator.gather(action_seq)
+        all_padding_seq = self._accelerator.gather(padding_seq)
+        self._collected_actions.append(
+            all_action_seq[torch.logical_not(all_padding_seq)]
+        )
+        if self._current_steps == self._kmeans_fit_steps:
+            logging.info("Fitting KMeans")
+            self._clustering_algo.fit(
+                torch.cat(self._collected_actions, dim=0).view(-1, self._act_dim)
+            )
+            self._have_fit_kmeans = True
+            self._cluster_centers = self._clustering_algo.bin_centers.float().to(
+                action_seq.device
+            )
+
+    def _predict(
+        self,
+        obs_seq: torch.Tensor,
+        goal_seq: Optional[torch.Tensor],
+        action_seq: Optional[torch.Tensor],
+        is_padded_action_seq: Optional[torch.Tensor],
+        predict_with_offset: bool = True,
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Dict[str, float]]:
+        batch_size, obs_T, _ = obs_seq.shape
+        _, action_T, _ = (
+            action_seq.shape if action_seq is not None else (None, None, None)
+        )
+        # Take the one that is not None.
+        actions_to_predict = action_T or obs_T
+        if self._num_extra_predicted_actions:
+            actions_to_predict += self._num_extra_predicted_actions
+        # Now, figure out if we should pad the obs seq.
+        if obs_T < actions_to_predict:
+            # We need to pad the obs seq.
+            pad_size = actions_to_predict - obs_T
+            padded_obs_seq = torch.cat(
+                [
+                    obs_seq,
+                    einops.repeat(
+                        self._obs_padding, "D -> N T D", N=batch_size, T=pad_size
+                    ),
+                ],
+                dim=1,
+            )
+        else:
+            padded_obs_seq = obs_seq
+        # Assume dimensions are N T D for N sequences of T timesteps with dimension D.
+        if self._cbet_method == self.GOAL_SPEC.unconditional:
+            gpt_input = padded_obs_seq
+        elif self._cbet_method == self.GOAL_SPEC.concat:
+            gpt_input = torch.cat([goal_seq, padded_obs_seq], dim=1)
+        elif self._cbet_method == self.GOAL_SPEC.stack:
+            gpt_input = torch.cat([goal_seq, padded_obs_seq], dim=-1)
+        else:
+            raise NotImplementedError
+
+        gpt_output = self._gpt_model(gpt_input)
+        if self._cbet_method == self.GOAL_SPEC.concat:
+            # Chop off the goal encodings.
+            gpt_output = gpt_output[:, goal_seq.size(1) :, :]
+        cbet_preds = self._map_to_cbet_preds(gpt_output)
+        cbet_logits, cbet_offsets = torch.split(
+            cbet_preds, [self._K, self._K * self._act_dim], dim=-1
+        )
+        cbet_offsets = einops.rearrange(cbet_offsets, "N T (K A) -> N T K A", K=self._K)
+
+        cbet_probs = torch.softmax(cbet_logits, dim=-1)
+        N, T, choices = cbet_probs.shape
+        # Sample from the multinomial distribution, one per row.
+        sampled_centers = einops.rearrange(
+            torch.multinomial(cbet_probs.view(-1, choices), num_samples=1),
+            "(N T) 1 -> N T 1",
+            N=N,
+        )
+        flattened_cbet_offsets = einops.rearrange(cbet_offsets, "N T K A -> (N T) K A")
+        sampled_offsets = flattened_cbet_offsets[
+            torch.arange(flattened_cbet_offsets.shape[0]), sampled_centers.flatten()
+        ].view(N, T, self._act_dim)
+        centers = self._cluster_centers[sampled_centers.flatten()].view(
+            N, T, self._act_dim
+        )
+        a_hat = centers
+        if predict_with_offset:
+            a_hat += sampled_offsets
+        if action_seq is None:
+            return a_hat, None, {}
+        # We are in training, so figure out the loss for the actions.
+        # First, we need to find the closest cluster center for each action.
+        action_bins = self._find_closest_cluster(action_seq)
+        true_offsets = action_seq - self._cluster_centers[action_bins]
+        predicted_offsets = flattened_cbet_offsets[
+            torch.arange(flattened_cbet_offsets.shape[0]), action_bins.flatten()
+        ].view(N, T, self._act_dim)
+        # Now we can compute the loss.
+        offset_loss = self._offset_criterion(predicted_offsets, true_offsets)
+        cbet_loss = self._criterion(
+            einops.rearrange(cbet_logits, "N T D -> (N T) D"),
+            einops.rearrange(action_bins, "N T -> (N T)"),
+        )
+        # Now, use the padding mask to mask out the loss.
+        if is_padded_action_seq is not None:
+            cbet_loss *= ~is_padded_action_seq.view(-1)
+            offset_loss *= ~is_padded_action_seq.unsqueeze(-1)
+        cbet_loss, offset_loss = cbet_loss.mean(), offset_loss.mean()
+        loss = cbet_loss + self._offset_loss_multiplier * offset_loss
+        action_mse = F.mse_loss(a_hat, action_seq, reduction="none")
+        action_l1 = F.l1_loss(a_hat, action_seq, reduction="none")
+        norm = torch.norm(action_seq, p=2, dim=-1, keepdim=True) + 1e-9
+        normalized_mse = (action_mse / norm).mean()
+        if self._current_steps < self._kmeans_fit_steps:
+            loss = loss.detach() + (loss * 0.0)
+
+        loss_dict = {
+            "classification_loss": cbet_loss.detach().cpu().item(),
+            "offset_loss": offset_loss.detach().cpu().item(),
+            "loss": loss.detach().cpu().item(),
+            "L2_loss": action_mse.mean().detach().cpu().item(),
+            "L2_loss_normalized": normalized_mse.mean().detach().cpu().item(),
+            "L1_loss": action_l1.mean().detach().cpu().item(),
+        }
+        return a_hat, loss, loss_dict
+
+    def _find_closest_cluster(self, action_seq: torch.Tensor) -> torch.Tensor:
+        N, T, _ = action_seq.shape
+        flattened_actions = einops.rearrange(action_seq, "N T A -> (N T) A")
+        cluster_center_distance = torch.sum(
+            (flattened_actions[:, None, :] - self._cluster_centers[None, :, :]) ** 2,
+            dim=2,
+        )  # (N T) K A -> (N T) K
+        closest_cluster_center = torch.argmin(cluster_center_distance, dim=1)  # (N T)
+        discretized_action = einops.rearrange(
+            closest_cluster_center, "(N T) -> N T", N=N, T=T
+        )
+        return discretized_action
+
+    def configure_optimizers(self, weight_decay, learning_rate, betas):
+        optimizer = self._gpt_model.configure_optimizers(
+            weight_decay=weight_decay,
+            learning_rate=learning_rate,
+            betas=betas,
+        )
+        optimizer.add_param_group({"params": self._map_to_cbet_preds.parameters()})
+        return optimizer
+
+
+class FocalLoss(nn.Module):
+    def __init__(self, gamma: float = 0, reduction: str = "mean"):
+        super(FocalLoss, self).__init__()
+        self.gamma = gamma
+        if reduction not in ("mean", "sum", "none"):
+            raise NotImplementedError
+        self.reduction = reduction
+
+    def forward(self, input, target):
+        logpt = F.log_softmax(input, dim=-1)
+        logpt = logpt.gather(1, target.view(-1, 1)).view(-1)
+        pt = logpt.exp()
+
+        loss = -1 * (1 - pt) ** self.gamma * logpt
+        if self.reduction == "mean":
+            return loss.mean()
+        elif self.reduction == "sum":
+            return loss.sum()
+        else:
+            return loss

baselines/rum/models/bet/gpt.py

@@ -0,0 +1,277 @@
+"""
+An adaptation of Andrej Karpathy's nanoGPT implementation in PyTorch.
+Original source: https://github.com/karpathy/nanoGPT
+
+Original License:
+MIT License
+
+Copyright (c) 2022 Andrej Karpathy
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+Original comments:
+Full definition of a GPT Language Model, all of it in this single file.
+References:
+1) the official GPT-2 TensorFlow implementation released by OpenAI:
+https://github.com/openai/gpt-2/blob/master/src/model.py
+2) huggingface/transformers PyTorch implementation:
+https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
+"""
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+# @torch.jit.script # good to enable when not using torch.compile, disable when using (our default)
+def new_gelu(x):
+    """
+    Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT).
+    Reference: Gaussian Error Linear Units (GELU) paper: https://arxiv.org/abs/1606.08415
+    """
+    return (
+        0.5
+        * x
+        * (
+            1.0
+            + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))
+        )
+    )
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        # causal mask to ensure that attention is only applied to the left in the input sequence
+        self.register_buffer(
+            "bias",
+            torch.tril(torch.ones(config.block_size, config.block_size)).view(
+                1, 1, config.block_size, config.block_size
+            ),
+        )
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+
+    def forward(self, x):
+        # batch size, sequence length, embedding dimensionality (n_embd)
+        B, T, C = x.size()
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(
+            1, 2
+        )  # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(
+            1, 2
+        )  # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(
+            1, 2
+        )  # (B, nh, T, hs)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+        att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float("-inf"))
+        att = F.softmax(att, dim=-1)
+        att = self.attn_dropout(att)
+        y = att @ v  # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = (
+            y.transpose(1, 2).contiguous().view(B, T, C)
+        )  # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = new_gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    input_dim: int = 256
+    output_dim: int = 256
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.1
+
+
+class GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.input_dim is not None
+        assert config.output_dim is not None
+        assert config.block_size is not None
+        self.config = config
+
+        self.transformer = nn.ModuleDict(
+            dict(
+                wte=nn.Linear(config.input_dim, config.n_embd),
+                wpe=nn.Embedding(config.block_size, config.n_embd),
+                drop=nn.Dropout(config.dropout),
+                h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+                ln_f=nn.LayerNorm(config.n_embd),
+            )
+        )
+        self.lm_head = nn.Linear(config.n_embd, config.output_dim, bias=False)
+        # init all weights, and apply a special scaled init to the residual projections, per GPT-2 paper
+        self.apply(self._init_weights)
+        for pn, p in self.named_parameters():
+            if pn.endswith("c_proj.weight"):
+                torch.nn.init.normal_(
+                    p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer)
+                )
+
+        # report number of parameters
+        n_params = sum(p.numel() for p in self.parameters())
+
+    def forward(self, input, targets=None):
+        device = input.device
+        b, t, d = input.size()
+        assert (
+            t <= self.config.block_size
+        ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(
+            0
+        )  # shape (1, t)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(
+            input
+        )  # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(
+            pos
+        )  # position embeddings of shape (1, t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+        logits = self.lm_head(x)
+        return logits
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            torch.nn.init.zeros_(module.bias)
+            torch.nn.init.ones_(module.weight)
+
+    def crop_block_size(self, block_size):
+        assert block_size <= self.config.block_size
+        self.config.block_size = block_size
+        self.transformer.wpe.weight = nn.Parameter(
+            self.transformer.wpe.weight[:block_size]
+        )
+        for block in self.transformer.h:
+            block.attn.bias = block.attn.bias[:, :, :block_size, :block_size]
+
+    def configure_optimizers(self, weight_decay, learning_rate, betas):
+        """
+        This long function is unfortunately doing something very simple and is being very defensive:
+        We are separating out all parameters of the model into two buckets: those that will experience
+        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).
+        We are then returning the PyTorch optimizer object.
+        """
+
+        # separate out all parameters to those that will and won't experience regularizing weight decay
+        decay = set()
+        no_decay = set()
+        whitelist_weight_modules = (torch.nn.Linear,)
+        blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
+        for mn, m in self.named_modules():
+            for pn, p in m.named_parameters():
+                fpn = "%s.%s" % (mn, pn) if mn else pn  # full param name
+                if pn.endswith("bias"):
+                    # all biases will not be decayed
+                    no_decay.add(fpn)
+                elif pn.endswith("weight") and isinstance(m, whitelist_weight_modules):
+                    # weights of whitelist modules will be weight decayed
+                    decay.add(fpn)
+                elif pn.endswith("weight") and isinstance(m, blacklist_weight_modules):
+                    # weights of blacklist modules will NOT be weight decayed
+                    no_decay.add(fpn)
+
+        # validate that we considered every parameter
+        param_dict = {pn: p for pn, p in self.named_parameters()}
+        inter_params = decay & no_decay
+        union_params = decay | no_decay
+        assert len(inter_params) == 0, (
+            "parameters %s made it into both decay/no_decay sets!"
+            % (str(inter_params),)
+        )
+        assert len(param_dict.keys() - union_params) == 0, (
+            "parameters %s were not separated into either decay/no_decay set!"
+            % (str(param_dict.keys() - union_params),)
+        )
+
+        # create the pytorch optimizer object
+        optim_groups = [
+            {
+                "params": [param_dict[pn] for pn in sorted(list(decay))],
+                "weight_decay": weight_decay,
+            },
+            {
+                "params": [param_dict[pn] for pn in sorted(list(no_decay))],
+                "weight_decay": 0.0,
+            },
+        ]
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas)
+        return optimizer