chore: import upstream snapshot with attribution

example/agents/helpers.py

@@ -0,0 +1,116 @@
+# Copyright 2022 Twitter, Inc and Zhendong Wang.
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+import time
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class SinusoidalPosEmb(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = x[:, None] * emb[None, :]
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+#-----------------------------------------------------------------------------#
+#---------------------------------- sampling ---------------------------------#
+#-----------------------------------------------------------------------------#
+
+
+def extract(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+
+def cosine_beta_schedule(timesteps, s=0.008, dtype=torch.float32):
+    """
+    cosine schedule
+    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
+    """
+    steps = timesteps + 1
+    x = np.linspace(0, steps, steps)
+    alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    betas_clipped = np.clip(betas, a_min=0, a_max=0.999)
+    return torch.tensor(betas_clipped, dtype=dtype)
+
+
+def linear_beta_schedule(timesteps, beta_start=1e-4, beta_end=2e-2, dtype=torch.float32):
+    betas = np.linspace(
+        beta_start, beta_end, timesteps
+    )
+    return torch.tensor(betas, dtype=dtype)
+
+
+def vp_beta_schedule(timesteps, dtype=torch.float32):
+    t = np.arange(1, timesteps + 1)
+    T = timesteps
+    b_max = 10.
+    b_min = 0.1
+    alpha = np.exp(-b_min / T - 0.5 * (b_max - b_min) * (2 * t - 1) / T ** 2)
+    betas = 1 - alpha
+    return torch.tensor(betas, dtype=dtype)
+
+#-----------------------------------------------------------------------------#
+#---------------------------------- losses -----------------------------------#
+#-----------------------------------------------------------------------------#
+
+class WeightedLoss(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, pred, targ, weights=1.0):
+        '''
+            pred, targ : tensor [ batch_size x action_dim ]
+        '''
+        loss = self._loss(pred, targ)
+        weighted_loss = (loss * weights).mean()
+        return weighted_loss
+
+class WeightedL1(WeightedLoss):
+
+    def _loss(self, pred, targ):
+        return torch.abs(pred - targ)
+
+class WeightedL2(WeightedLoss):
+
+    def _loss(self, pred, targ):
+        return F.mse_loss(pred, targ, reduction='none')
+
+
+Losses = {
+    'l1': WeightedL1,
+    'l2': WeightedL2,
+}
+
+
+class EMA():
+    '''
+        empirical moving average
+    '''
+    def __init__(self, beta):
+        super().__init__()
+        self.beta = beta
+
+    def update_model_average(self, ma_model, current_model):
+        for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
+            old_weight, up_weight = ma_params.data, current_params.data
+            ma_params.data = self.update_average(old_weight, up_weight)
+
+    def update_average(self, old, new):
+        if old is None:
+            return new
+        return old * self.beta + (1 - self.beta) * new