Upload 10 files

README.md

@@ -0,0 +1,62 @@
+## Policy Representation via Diffusion Probability Model for Reinforcement Learning
+
+**Policy Representation via Diffusion Probability Model for Reinforcement Learning**<br>
+Anonymous <br>
+
+Abstract: *Popular reinforcement learning (RL) algorithms tend to produce a unimodal policy distribution, which weakens the expressiveness of complicated policy and decays the ability of exploration. The diffusion probability model is powerful to learn complicated multimodal distributions, which has shown promising and potential applications to RL. In this paper, we formally build a theoretical foundation of policy representation via the diffusion probability model and provide practical implementations of diffusion policy for online model-free RL. Concretely, we character diffusion policy as a stochastic process, which is a new approach to representing a policy. Then we present a convergence guarantee for diffusion policy, which provides a theory to understand the multimodality of diffusion policy. Furthermore, we propose the DIPO which is an implementation for model-free online RL with \textbf{DI}ffusion \textbf{PO}licy. To the best of our knowledge, DIPO is the first algorithm to solve model-free online RL problems with the diffusion model. Finally, extensive empirical results show the effectiveness and superiority of DIPO on the standard continuous control MoJoCo benchmark.*
+
+## Experiments
+
+### Requirements
+Installations of [PyTorch](https://pytorch.org/) and [MuJoCo](https://github.com/deepmind/mujoco) are needed. 
+A suitable [conda](https://conda.io) environment named `DIPO` can be created and activated with:
+```.bash
+conda create DIPO
+conda activate DIPO
+```
+To get started, install the additionally required python packages into you environment.
+```.bash
+pip install -r requirements.txt
+```
+
+### Running
+Running experiments based our code could be quite easy, so below we use `Hopper-v3` task as an example. 
+
+```.bash
+python main.py --env_name Hopper-v3 --num_steps 1000000 --n_timesteps 100 --cuda 0 --seed 0
+```
+
+
+### Hyperparameters
+Hyperparameters for DIPO have been shown as follow for easily reproducing our reported results.
+
+#### Hyper-parameters for algorithms
+| Hyperparameter | DIPO | SAC | TD3 | PPO |
+| -------------- | ---- | --- | --- | --- |
+| No. of hidden layers | 2 | 2 | 2 | 2 |
+| No. of hidden nodes | 256 | 256  | 256  | 256  |
+| Activation | mish | relu | relu | tanh |
+| Batch size | 256 | 256 | 256 | 256 |
+| Discount for reward $\gamma$ | 0.99 | 0.99 | 0.99 | 0.99 |
+| Target smoothing coefficient $\tau$ | 0.005 | 0.005 | 0.005 | 0.005 |
+| Learning rate for actor | $3 × 10^{-4}$ | $3 × 10^{-4}$ | $3 × 10^{-4}$ | $7 × 10^{-4}$ |
+| Learning rate for actor | $3 × 10^{-4}$ | $3 × 10^{-4}$ | $3 × 10^{-4}$ | $7 × 10^{-4}$ |
+| Actor Critic grad norm | 2 | N/A | N/A | 0.5 |
+| Memeroy size | $1 × 10^6$ | $1 × 10^6$ | $1 × 10^6$ | $1 × 10^6$ |
+| Entropy coefficient | N/A | 0.2 | N/A | 0.01 |
+| Value loss coefficient | N/A | N/A | N/A | 0.5 |
+| Exploration noise | N/A | N/A | $\mathcal{N}$(0, 0.1) | N/A |
+| Policy noise | N/A | N/A | $\mathcal{N}$(0, 0.2) | N/A |
+| Noise clip | N/A | N/A | 0.5 | N/A |
+| Use gae | N/A | N/A | N/A | True |
+
+#### Hyper-parameters for MuJoCo.(DIPO)
+| Hyperparameter | Hopper-v3 | Walker2d-v3 | Ant-v3 | HalfCheetah-v3 | Humanoid-v3 |
+| --- | --- | --- | --- | --- | --- |
+| Learning rate for action | 0.03 | 0.03 | 0.03 | 0.03 | 0.03 |
+| Actor Critic grad norm | 1 | 2 | 0.8 | 2 | 2 |
+| Action grad norm ratio | 0.3 | 0.08 | 0.1 | 0.08 | 0.1 |
+| Action gradient steps | 20 | 20 | 20 | 40 | 20 |
+| Diffusion inference timesteps | 100 | 100 | 100 | 100 | 100 |
+| Diffusion beta schedule | cosine | cosine | cosine | cosine | cosine |
+| Update actor target every | 1 | 1 | 1 | 2 | 1 |

agent/DiPo.py

@@ -0,0 +1,173 @@
+import copy
+import numpy as np
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
+from agent.model import MLP, Critic
+from agent.diffusion import Diffusion
+from agent.vae import VAE
+from agent.helpers import EMA
+
+
+class DiPo(object):
+    def __init__(self,
+                 args,
+                 state_dim,
+                 action_space,
+                 memory,
+                 diffusion_memory,
+                 device,
+                 ):
+        action_dim = np.prod(action_space.shape)
+
+        self.policy_type = args.policy_type
+        if self.policy_type == 'Diffusion':
+            self.actor = Diffusion(state_dim=state_dim, action_dim=action_dim, noise_ratio=args.noise_ratio,
+                                   beta_schedule=args.beta_schedule, n_timesteps=args.n_timesteps).to(device)
+        elif self.policy_type == 'VAE':
+            self.actor = VAE(state_dim=state_dim, action_dim=action_dim, device=device).to(device)
+        else:
+            self.actor = MLP(state_dim=state_dim, action_dim=action_dim).to(device)
+
+        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=args.diffusion_lr, eps=1e-5)
+
+        self.memory = memory
+        self.diffusion_memory = diffusion_memory
+        self.action_gradient_steps = args.action_gradient_steps
+
+        self.action_grad_norm = action_dim * args.ratio
+        self.ac_grad_norm = args.ac_grad_norm
+
+        self.step = 0
+        self.tau = args.tau
+        self.actor_target = copy.deepcopy(self.actor)
+        self.update_actor_target_every = args.update_actor_target_every
+
+        self.critic = Critic(state_dim, action_dim).to(device)
+        self.critic_target = copy.deepcopy(self.critic)
+        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=args.critic_lr, eps=1e-5)
+
+        self.action_dim = action_dim
+
+        self.action_lr = args.action_lr
+
+        self.device = device
+
+        if action_space is None:
+            self.action_scale = 1.
+            self.action_bias = 0.
+        else:
+            self.action_scale = (action_space.high - action_space.low) / 2.
+            self.action_bias = (action_space.high + action_space.low) / 2.
+
+    def append_memory(self, state, action, reward, next_state, mask):
+        action = (action - self.action_bias) / self.action_scale
+        
+        self.memory.append(state, action, reward, next_state, mask)
+        self.diffusion_memory.append(state, action)
+
+    def sample_action(self, state, eval=False):
+        state = torch.FloatTensor(state.reshape(1, -1)).to(self.device)
+
+        action = self.actor(state, eval).cpu().data.numpy().flatten()
+        action = action.clip(-1, 1)
+        action = action * self.action_scale + self.action_bias
+        return action
+
+    def action_gradient(self, batch_size, log_writer):
+        states, best_actions, idxs = self.diffusion_memory.sample(batch_size)
+
+        actions_optim = torch.optim.Adam([best_actions], lr=self.action_lr, eps=1e-5)
+
+
+        for i in range(self.action_gradient_steps):
+            best_actions.requires_grad_(True)
+            q1, q2 = self.critic(states, best_actions)
+            loss = -torch.min(q1, q2)
+
+            actions_optim.zero_grad()
+
+            loss.backward(torch.ones_like(loss))
+            if self.action_grad_norm > 0:
+                actions_grad_norms = nn.utils.clip_grad_norm_([best_actions], max_norm=self.action_grad_norm, norm_type=2)
+
+            actions_optim.step()
+
+            best_actions.requires_grad_(False)
+            best_actions.clamp_(-1., 1.)
+
+        # if self.step % 10 == 0:
+        #     log_writer.add_scalar('Action Grad Norm', actions_grad_norms.max().item(), self.step)
+
+        best_actions = best_actions.detach()
+
+        self.diffusion_memory.replace(idxs, best_actions.cpu().numpy())
+
+        return states, best_actions
+
+    def train(self, iterations, batch_size=256, log_writer=None):
+        for _ in range(iterations):
+            # Sample replay buffer / batch
+            states, actions, rewards, next_states, masks = self.memory.sample(batch_size)
+
+            """ Q Training """
+            current_q1, current_q2 = self.critic(states, actions)
+
+            next_actions = self.actor_target(next_states, eval=False)
+            target_q1, target_q2 = self.critic_target(next_states, next_actions)
+            target_q = torch.min(target_q1, target_q2)
+
+            target_q = (rewards + masks * target_q).detach()
+
+            critic_loss = F.mse_loss(current_q1, target_q) + F.mse_loss(current_q2, target_q)
+
+            self.critic_optimizer.zero_grad()
+            critic_loss.backward()
+            if self.ac_grad_norm > 0:
+                critic_grad_norms = nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=self.ac_grad_norm, norm_type=2)
+                # if self.step % 10 == 0:
+                #     log_writer.add_scalar('Critic Grad Norm', critic_grad_norms.max().item(), self.step)
+            self.critic_optimizer.step()
+
+            """ Policy Training """
+            states, best_actions = self.action_gradient(batch_size, log_writer)
+
+            actor_loss = self.actor.loss(best_actions, states)
+
+            self.actor_optimizer.zero_grad()
+            actor_loss.backward()
+            if self.ac_grad_norm > 0:
+                actor_grad_norms = nn.utils.clip_grad_norm_(self.actor.parameters(), max_norm=self.ac_grad_norm, norm_type=2)
+                # if self.step % 10 == 0:
+                #     log_writer.add_scalar('Actor Grad Norm', actor_grad_norms.max().item(), self.step)
+            self.actor_optimizer.step()
+
+            """ Step Target network """
+            for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
+                target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+
+            if self.step % self.update_actor_target_every == 0:
+                for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):
+                    target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+
+            self.step += 1
+
+    def save_model(self, dir, id=None):
+        if id is not None:
+            torch.save(self.actor.state_dict(), f'{dir}/actor_{id}.pth')
+            torch.save(self.critic.state_dict(), f'{dir}/critic_{id}.pth')
+        else:
+            torch.save(self.actor.state_dict(), f'{dir}/actor.pth')
+            torch.save(self.critic.state_dict(), f'{dir}/critic.pth')
+
+    def load_model(self, dir, id=None):
+        if id is not None:
+            self.actor.load_state_dict(torch.load(f'{dir}/actor_{id}.pth'))
+            self.critic.load_state_dict(torch.load(f'{dir}/critic_{id}.pth'))
+        else:
+            self.actor.load_state_dict(torch.load(f'{dir}/actor.pth'))
+            self.critic.load_state_dict(torch.load(f'{dir}/critic.pth'))
+

agent/diffusion.py

@@ -0,0 +1,178 @@
+import copy
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+from agent.helpers import (cosine_beta_schedule,
+                            linear_beta_schedule,
+                            vp_beta_schedule,
+                            extract,
+                            Losses)
+
+from agent.model import Model
+
+
+class Diffusion(nn.Module):
+    def __init__(self, state_dim, action_dim, noise_ratio,
+                 beta_schedule='vp', n_timesteps=1000,
+                 loss_type='l2', clip_denoised=True, predict_epsilon=True):
+        super(Diffusion, self).__init__()
+
+        self.state_dim = state_dim
+        self.action_dim = action_dim
+        self.model = Model(state_dim, action_dim)
+
+        self.max_noise_ratio = noise_ratio
+        self.noise_ratio = noise_ratio
+
+        if beta_schedule == 'linear':
+            betas = linear_beta_schedule(n_timesteps)
+        elif beta_schedule == 'cosine':
+            betas = cosine_beta_schedule(n_timesteps)
+        elif beta_schedule == 'vp':
+            betas = vp_beta_schedule(n_timesteps)
+
+        alphas = 1. - betas
+        alphas_cumprod = torch.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]])
+
+        self.n_timesteps = int(n_timesteps)
+        self.clip_denoised = clip_denoised
+        self.predict_epsilon = predict_epsilon
+
+        self.register_buffer('betas', betas)
+        self.register_buffer('alphas_cumprod', alphas_cumprod)
+        self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod))
+        self.register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod))
+        self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1))
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
+        self.register_buffer('posterior_variance', posterior_variance)
+
+        ## log calculation clipped because the posterior variance
+        ## is 0 at the beginning of the diffusion chain
+        self.register_buffer('posterior_log_variance_clipped',
+                             torch.log(torch.clamp(posterior_variance, min=1e-20)))
+        self.register_buffer('posterior_mean_coef1',
+                             betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
+        self.register_buffer('posterior_mean_coef2',
+                             (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))
+
+        self.loss_fn = Losses[loss_type]()
+
+    # ------------------------------------------ sampling ------------------------------------------#
+
+    def predict_start_from_noise(self, x_t, t, noise):
+        '''
+            if self.predict_epsilon, model output is (scaled) noise;
+            otherwise, model predicts x0 directly
+        '''
+        if self.predict_epsilon:
+            return (
+                    extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
+                    extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
+            )
+        else:
+            return noise
+
+    def q_posterior(self, x_start, x_t, t):
+        posterior_mean = (
+                extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
+                extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(self, x, t, s):
+        x_recon = self.predict_start_from_noise(x, t=t, noise=self.model(x, t, s))
+
+        if self.clip_denoised:
+            x_recon.clamp_(-1., 1.)
+        else:
+            assert RuntimeError()
+
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
+        return model_mean, posterior_variance, posterior_log_variance
+
+    @torch.no_grad()
+    def p_sample(self, x, t, s):
+        b, *_, device = *x.shape, x.device
+
+        model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, s=s)
+
+        noise = torch.randn_like(x)
+        # no noise when t == 0
+        nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+
+        return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise * self.noise_ratio
+
+
+    @torch.no_grad()
+    def p_sample_loop(self, state, shape):
+        device = self.betas.device
+
+        batch_size = shape[0]
+        x = torch.randn(shape, device=device)
+
+        for i in reversed(range(0, self.n_timesteps)):
+            timesteps = torch.full((batch_size,), i, device=device, dtype=torch.long)
+            x = self.p_sample(x, timesteps, state)
+
+        return x
+
+    @torch.no_grad()
+    def sample(self, state, eval=False):
+        self.noise_ratio = 0 if eval else self.max_noise_ratio
+        
+        batch_size = state.shape[0]
+        shape = (batch_size, self.action_dim)
+        action = self.p_sample_loop(state, shape)
+        return action.clamp_(-1., 1.)
+
+    # ------------------------------------------ training ------------------------------------------#
+
+    def q_sample(self, x_start, t, noise=None):
+        if noise is None:
+            noise = torch.randn_like(x_start)
+
+        sample = (
+                extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
+                extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+        )
+
+        return sample
+
+    def p_losses(self, x_start, state, t, weights=1.0):
+        noise = torch.randn_like(x_start)
+
+        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+
+        x_recon = self.model(x_noisy, t, state)
+
+        assert noise.shape == x_recon.shape
+
+        if self.predict_epsilon:
+            loss = self.loss_fn(x_recon, noise, weights)
+        else:
+            loss = self.loss_fn(x_recon, x_start, weights)
+
+        return loss
+
+
+    def loss(self, x, state, weights=1.0):
+        batch_size = len(x)
+        t = torch.randint(0, self.n_timesteps, (batch_size,), device=x.device).long()
+        return self.p_losses(x, state, t, weights)
+
+    def forward(self, state, eval=False):
+        return self.sample(state, eval)
+

agent/helpers.py

@@ -0,0 +1,130 @@
+import math
+import time
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def init_weights(m):
+    def truncated_normal_init(t, mean=0.0, std=0.01):
+        torch.nn.init.normal_(t, mean=mean, std=std)
+        while True:
+            cond = torch.logical_or(t < mean - 2 * std, t > mean + 2 * std)
+            if not torch.sum(cond):
+                break
+            t = torch.where(cond, torch.nn.init.normal_(torch.ones(t.shape), mean=mean, std=std), t)
+        return t
+
+    if type(m) == nn.Linear:
+        input_dim = m.in_features
+        truncated_normal_init(m.weight, std=1 / (2 * np.sqrt(input_dim)))
+        m.bias.data.fill_(0.0)
+
+
+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

agent/model.py

@@ -0,0 +1,100 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from agent.helpers import SinusoidalPosEmb, init_weights
+
+
+class Critic(nn.Module):
+    def __init__(self, state_dim, action_dim, hidden_dim=256):
+        super(Critic, self).__init__()
+        self.q1_model = nn.Sequential(nn.Linear(state_dim + action_dim, hidden_dim),
+                                      nn.Mish(),
+                                      nn.Linear(hidden_dim, hidden_dim),
+                                      nn.Mish(),
+                                      nn.Linear(hidden_dim, hidden_dim),
+                                      nn.Mish(),
+                                      nn.Linear(hidden_dim, 1))
+
+        self.q2_model = nn.Sequential(nn.Linear(state_dim + action_dim, hidden_dim),
+                                      nn.Mish(),
+                                      nn.Linear(hidden_dim, hidden_dim),
+                                      nn.Mish(),
+                                      nn.Linear(hidden_dim, hidden_dim),
+                                      nn.Mish(),
+                                      nn.Linear(hidden_dim, 1))
+
+        self.apply(init_weights)
+
+    def forward(self, state, action):
+        x = torch.cat([state, action], dim=-1)
+        return self.q1_model(x), self.q2_model(x)
+
+    def q1(self, state, action):
+        x = torch.cat([state, action], dim=-1)
+        return self.q1_model(x)
+
+    def q_min(self, state, action):
+        q1, q2 = self.forward(state, action)
+        return torch.min(q1, q2)
+
+
+class Model(nn.Module):
+    def __init__(self, state_dim, action_dim, hidden_size=256, time_dim=32):
+        super(Model, self).__init__()
+
+        self.time_mlp = nn.Sequential(
+            SinusoidalPosEmb(time_dim),
+            nn.Linear(time_dim, hidden_size),
+            nn.Mish(),
+            nn.Linear(hidden_size, time_dim),
+        )
+
+        input_dim = state_dim + action_dim + time_dim
+        self.layer = nn.Sequential(nn.Linear(input_dim, hidden_size),
+                                       nn.Mish(),
+                                       nn.Linear(hidden_size, hidden_size),
+                                       nn.Mish(),
+                                       nn.Linear(hidden_size, hidden_size),
+                                       nn.Mish(),
+                                       nn.Linear(hidden_size, action_dim))
+        self.apply(init_weights)
+        
+
+    def forward(self, x, time, state):
+
+        t = self.time_mlp(time)
+        out = torch.cat([x, t, state], dim=-1)
+        out = self.layer(out)
+
+        return out
+
+
+class MLP(nn.Module):
+    def __init__(self, state_dim, action_dim, hidden_size=256):
+        super(MLP, self).__init__()
+
+        input_dim = state_dim
+        self.mid_layer = nn.Sequential(nn.Linear(input_dim, hidden_size),
+                                       nn.Mish(),
+                                       nn.Linear(hidden_size, hidden_size),
+                                       nn.Mish(),
+                                       nn.Linear(hidden_size, hidden_size),
+                                       nn.Mish())
+        
+        self.final_layer = nn.Linear(hidden_size, action_dim)
+
+        self.apply(init_weights)
+
+    def forward(self, state, eval=False):
+        out = self.mid_layer(state)
+        out = self.final_layer(out)
+
+        if not eval:
+            out += torch.randn_like(out) * 0.1
+
+        return out
+
+    def loss(self, action, state):
+        return F.mse_loss(self.forward(state), action, reduction='mean')

agent/replay_memory.py

@@ -0,0 +1,79 @@
+import numpy as np
+import torch
+
+
+class ReplayMemory():
+    """Buffer to store environment transitions."""
+    def __init__(self, state_dim, action_dim, capacity, device):
+        self.capacity = int(capacity)
+        self.device = device
+
+        self.states = np.empty((self.capacity, int(state_dim)), dtype=np.float32)
+        self.actions = np.empty((self.capacity, int(action_dim)), dtype=np.float32)
+        self.rewards = np.empty((self.capacity, 1), dtype=np.float32)
+        self.next_states = np.empty((self.capacity, int(state_dim)), dtype=np.float32)
+        self.masks = np.empty((self.capacity, 1), dtype=np.float32)
+
+        self.idx = 0
+        self.full = False
+
+    def append(self, state, action, reward, next_state, mask):
+
+        np.copyto(self.states[self.idx], state)
+        np.copyto(self.actions[self.idx], action)
+        np.copyto(self.rewards[self.idx], reward)
+        np.copyto(self.next_states[self.idx], next_state)
+        np.copyto(self.masks[self.idx], mask)
+
+        self.idx = (self.idx + 1) % self.capacity
+        self.full = self.full or self.idx == 0
+
+    def sample(self, batch_size):
+        idxs = np.random.randint(
+            0, self.capacity if self.full else self.idx, size=batch_size
+        )
+
+        states = torch.as_tensor(self.states[idxs], device=self.device)
+        actions = torch.as_tensor(self.actions[idxs], device=self.device)
+        rewards = torch.as_tensor(self.rewards[idxs], device=self.device)
+        next_states = torch.as_tensor(self.next_states[idxs], device=self.device)
+        masks = torch.as_tensor(self.masks[idxs], device=self.device)
+
+        return states, actions, rewards, next_states, masks
+
+
+class DiffusionMemory():
+    """Buffer to store best actions."""
+    def __init__(self, state_dim, action_dim, capacity, device):
+        self.capacity = int(capacity)
+        self.device = device
+
+        self.states = np.empty((self.capacity, int(state_dim)), dtype=np.float32)
+        self.best_actions = np.empty((self.capacity, int(action_dim)), dtype=np.float32)
+
+        self.idx = 0
+        self.full = False
+
+    def append(self, state, action):
+
+        np.copyto(self.states[self.idx], state)
+        np.copyto(self.best_actions[self.idx], action)
+
+        self.idx = (self.idx + 1) % self.capacity
+        self.full = self.full or self.idx == 0
+
+    def sample(self, batch_size):
+        idxs = np.random.randint(
+            0, self.capacity if self.full else self.idx, size=batch_size
+        )
+
+        states = torch.as_tensor(self.states[idxs], device=self.device)
+        best_actions = torch.as_tensor(self.best_actions[idxs], device=self.device)
+
+        best_actions.requires_grad_(True)
+
+        return states, best_actions, idxs
+
+    def replace(self, idxs, best_actions):
+        np.copyto(self.best_actions[idxs], best_actions)
+    

agent/vae.py

@@ -0,0 +1,95 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from agent.helpers import init_weights
+
+
+class VAE(nn.Module):
+    def __init__(self, state_dim, action_dim, device, hidden_size=256) -> None:
+        super(VAE, self).__init__()
+
+        self.hidden_size = hidden_size
+        self.action_dim = action_dim
+
+        input_dim = state_dim + action_dim
+
+        self.encoder = nn.Sequential(nn.Linear(input_dim, hidden_size),
+                                     nn.Mish(),
+                                     nn.Linear(hidden_size, hidden_size),
+                                     nn.Mish(),
+                                     nn.Linear(hidden_size, hidden_size),
+                                     nn.Mish())
+
+        self.fc_mu = nn.Linear(hidden_size, hidden_size)
+        self.fc_var = nn.Linear(hidden_size, hidden_size)
+
+        self.decoder = nn.Sequential(nn.Linear(hidden_size + state_dim, hidden_size),
+                                     nn.Mish(),
+                                     nn.Linear(hidden_size, hidden_size),
+                                     nn.Mish(),
+                                     nn.Linear(hidden_size, hidden_size),
+                                     nn.Mish())
+
+        self.final_layer = nn.Sequential(nn.Linear(hidden_size, action_dim))
+
+        self.apply(init_weights)
+
+        self.device = device
+
+    def encode(self, action, state):
+        x = torch.cat([action, state], dim=-1)
+        result = self.encoder(x)
+        result = torch.flatten(result, start_dim=1)
+
+        # Split the result into mu and var components
+        # of the latent Gaussian distribution
+        mu = self.fc_mu(result)
+        log_var = self.fc_var(result)
+
+        return mu, log_var
+
+    def decode(self, z, state):
+        x = torch.cat([z, state], dim=-1)
+        result = self.decoder(x)
+        result = self.final_layer(result)
+        return result
+
+    def reparameterize(self, mu, logvar):
+        """
+        Will a single z be enough ti compute the expectation
+        for the loss??
+        :param mu: (Tensor) Mean of the latent Gaussian
+        :param logvar: (Tensor) Standard deviation of the latent Gaussian
+        :return:
+        """
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def loss(self, action, state):
+        mu, log_var = self.encode(action, state)
+        z = self.reparameterize(mu, log_var)
+        recons = self.decode(z, state)
+
+        kld_weight = 0.1  # Account for the minibatch samples from the dataset
+        recons_loss = F.mse_loss(recons, action)
+
+        kld_loss = torch.mean(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim=1), dim=0)
+
+        # print('recons_loss: ', recons_loss)
+        # print('kld_loss: ', kld_loss)
+
+        loss = recons_loss + kld_weight * kld_loss
+        return loss
+
+    def forward(self, state, eval=False):
+        batch_size = state.shape[0]
+        shape = (batch_size, self.hidden_size)
+
+        if eval:
+            z = torch.zeros(shape, device=self.device)
+        else:
+            z = torch.randn(shape, device=self.device)
+        samples = self.decode(z, state)
+
+        return samples.clamp(-1., 1.)

main.py

@@ -0,0 +1,163 @@
+import argparse
+import numpy as np
+import torch
+
+from agent.DiPo import DiPo
+from agent.replay_memory import ReplayMemory, DiffusionMemory
+
+from tensorboardX import SummaryWriter
+import gym
+import os
+
+
+def readParser():
+    parser = argparse.ArgumentParser(description='Diffusion Policy')
+    parser.add_argument('--env_name', default="Hopper-v3",
+                        help='Mujoco Gym environment (default: Hopper-v3)')
+    parser.add_argument('--seed', type=int, default=0, metavar='N',
+                        help='random seed (default: 0)')
+
+    parser.add_argument('--num_steps', type=int, default=1000000, metavar='N',
+                        help='env timesteps (default: 1000000)')
+
+    parser.add_argument('--batch_size', type=int, default=256, metavar='N',
+                        help='batch size (default: 256)')
+    parser.add_argument('--gamma', type=float, default=0.99, metavar='G',
+                        help='discount factor for reward (default: 0.99)')
+    parser.add_argument('--tau', type=float, default=0.005, metavar='G',
+                        help='target smoothing coefficient(τ) (default: 0.005)')
+    parser.add_argument('--update_actor_target_every', type=int, default=1, metavar='N',
+                        help='update actor target per iteration (default: 1)')
+
+    parser.add_argument("--policy_type", type=str, default="Diffusion", metavar='S',
+                        help="Diffusion, VAE or MLP")
+    parser.add_argument("--beta_schedule", type=str, default="cosine", metavar='S',
+                        help="linear, cosine or vp")
+    parser.add_argument('--n_timesteps', type=int, default=100, metavar='N',
+                        help='diffusion timesteps (default: 100)')
+    parser.add_argument('--diffusion_lr', type=float, default=0.0003, metavar='G',
+                        help='diffusion learning rate (default: 0.0003)')
+    parser.add_argument('--critic_lr', type=float, default=0.0003, metavar='G',
+                        help='critic learning rate (default: 0.0003)')
+    parser.add_argument('--action_lr', type=float, default=0.03, metavar='G',
+                        help='diffusion learning rate (default: 0.03)')
+    parser.add_argument('--noise_ratio', type=float, default=1.0, metavar='G',
+                        help='noise ratio in sample process (default: 1.0)')
+
+    parser.add_argument('--action_gradient_steps', type=int, default=20, metavar='N',
+                        help='action gradient steps (default: 20)')
+    parser.add_argument('--ratio', type=float, default=0.1, metavar='G',
+                        help='the ratio of action grad norm to action_dim (default: 0.1)')
+    parser.add_argument('--ac_grad_norm', type=float, default=2.0, metavar='G',
+                        help='actor and critic grad norm (default: 1.0)')
+
+    parser.add_argument('--cuda', default='cuda:0',
+                        help='run on CUDA (default: cuda:0)')
+
+    return parser.parse_args()
+
+
+def evaluate(env, agent, writer, steps):
+    episodes = 10
+    returns = np.zeros((episodes,), dtype=np.float32)
+
+    for i in range(episodes):
+        state = env.reset()
+        episode_reward = 0.
+        done = False
+        while not done:
+            action = agent.sample_action(state, eval=True)
+            next_state, reward, done, _ = env.step(action)
+            episode_reward += reward
+            state = next_state
+        returns[i] = episode_reward
+
+    mean_return = np.mean(returns)
+
+    writer.add_scalar(
+            'reward/test', mean_return, steps)
+    print('-' * 60)
+    print(f'Num steps: {steps:<5}  '
+              f'reward: {mean_return:<5.1f}')
+    print('-' * 60)
+
+
+def main(args=None):
+    if args is None:
+        args = readParser()
+
+    device = torch.device(args.cuda)
+
+    dir = "record"
+    # dir = "test"
+    log_dir = os.path.join(dir, f'{args.env_name}', f'policy_type={args.policy_type}', f'ratio={args.ratio}', f'seed={args.seed}')
+    writer = SummaryWriter(log_dir)
+
+    # Initial environment
+    env = gym.make(args.env_name)
+    state_size = int(np.prod(env.observation_space.shape))
+    action_size = int(np.prod(env.action_space.shape))
+    print(action_size)
+
+    # Set random seed
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    env.seed(args.seed)
+
+    memory_size = 1e6
+    num_steps = args.num_steps
+    start_steps = 10000
+    eval_interval = 10000
+    updates_per_step = 1
+    batch_size = args.batch_size
+    log_interval = 10
+    
+    memory = ReplayMemory(state_size, action_size, memory_size, device)
+    diffusion_memory = DiffusionMemory(state_size, action_size, memory_size, device)
+
+    agent = DiPo(args, state_size, env.action_space, memory, diffusion_memory, device)
+
+    steps = 0
+    episodes = 0
+
+    while steps < num_steps:
+        episode_reward = 0.
+        episode_steps = 0
+        done = False
+        state = env.reset()
+        episodes += 1
+        while not done:
+            if start_steps > steps:
+                action = env.action_space.sample()
+            else:
+                action = agent.sample_action(state, eval=False)
+            next_state, reward, done, _ = env.step(action)
+
+            mask = 0.0 if done else args.gamma
+
+            steps += 1
+            episode_steps += 1
+            episode_reward += reward
+
+            agent.append_memory(state, action, reward, next_state, mask)
+
+            if steps >= start_steps:
+                agent.train(updates_per_step, batch_size=batch_size, log_writer=writer)
+
+            if steps % eval_interval == 0:
+                evaluate(env, agent, writer, steps)
+                # self.save_models()
+                done =True
+
+            state = next_state
+
+        # if episodes % log_interval == 0:
+        #     writer.add_scalar('reward/train', episode_reward, steps)
+
+        print(f'episode: {episodes:<4}  '
+            f'episode steps: {episode_steps:<4}  '
+            f'reward: {episode_reward:<5.1f}')
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,103 @@
+absl-py==1.4.0
+ale-py==0.8.1
+asttokens==2.2.1
+attrs==22.2.0
+backcall==0.2.0
+box2d-py==2.3.5
+cachetools==5.3.0
+certifi==2022.12.7
+cffi==1.15.1
+charset-normalizer==3.1.0
+click==8.1.3
+cloudpickle==2.2.1
+cmake==3.26.0
+contourpy==1.0.7
+cycler==0.11.0
+Cython==0.29.33
+decorator==4.4.2
+docopt==0.6.2
+executing==1.2.0
+fasteners==0.18
+filelock==3.10.0
+fonttools==4.39.2
+glfw==2.5.7
+grpcio==1.51.3
+gym==0.21.0
+gym-notices==0.0.8
+h5py==3.8.0
+idna==3.4
+imageio==2.26.0
+imageio-ffmpeg==0.4.8
+importlib-metadata==4.13.0
+importlib-resources==5.12.0
+iniconfig==2.0.0
+ipython==8.11.0
+jedi==0.18.2
+Jinja2==3.1.2
+kiwisolver==1.4.4
+labmaze==1.0.6
+lit==15.0.7
+lxml==4.9.2
+lz4==4.3.2
+Markdown==3.4.1
+MarkupSafe==2.1.2
+matplotlib==3.7.1
+matplotlib-inline==0.1.6
+mjrl @ git+https://github.com/aravindr93/mjrl@3871d93763d3b49c4741e6daeaebbc605fe140dc
+moviepy==1.0.3
+mpmath==1.3.0
+mujoco==2.3.2
+mujoco-py==2.1.2.14
+networkx==3.0
+numpy==1.24.2
+oauthlib==3.2.2
+packaging==23.0
+pandas==1.5.3
+parso==0.8.3
+pexpect==4.8.0
+pickleshare==0.7.5
+Pillow==9.4.0
+pipreqs==0.4.13
+pluggy==1.0.0
+proglog==0.1.10
+prompt-toolkit==3.0.38
+protobuf==3.20.3
+ptyprocess==0.7.0
+pure-eval==0.2.2
+py==1.11.0
+pyasn1==0.4.8
+pyasn1-modules==0.2.8
+pybullet==3.2.5
+pycparser==2.21
+pygame==2.1.0
+Pygments==2.14.0
+PyOpenGL==3.1.6
+pyparsing==3.0.9
+pytest==7.0.1
+python-dateutil==2.8.2
+pytz==2022.7.1
+requests==2.28.2
+requests-oauthlib==1.3.1
+rsa==4.9
+scipy==1.10.1
+six==1.16.0
+stable-baselines3==1.7.0
+stack-data==0.6.2
+swig==4.1.1
+sympy==1.11.1
+tensorboard==2.12.0
+tensorboard-data-server==0.7.0
+tensorboard-plugin-wit==1.8.1
+tensorboardX==2.6
+termcolor==2.3.0
+tomli==2.0.1
+torch==2.0.0
+tqdm==4.65.0
+traitlets==5.9.0
+triton==2.0.0
+typing_extensions==4.5.0
+urllib3==1.26.15
+wcwidth==0.2.6
+Werkzeug==2.2.3
+yarg==0.1.9
+zipp==3.15.0

run_dipo

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+# Script to reproduce results
+
+envs=(Hopper-v3 Walker2d-v3 Ant-v3 HalfCheetah-v3 Humanoid-v3)
+steps=(1000000 1000000 3000000 3000000 10000000)
+cnt=0
+i=3
+n_timesteps=100
+for ((j=0;j<5;j+=1))
+do
+	nohup python -u main.py \
+	--env_name ${envs[i]} \
+	--num_steps 1000000 \
+	--policy_type 'MLP' \
+	--beta_schedule 'cosine' \
+	--n_timesteps  ${n_timesteps}\
+	--ratio 0.08 \
+	--ac_grad_norm 2 \
+	--action_gradient_steps 40 \
+	--update_actor_target_every 2 \
+	--seed $j \
+	--cuda "cuda:${cnt}" \
+	> "log/MLP-a_steps=40-%2-${envs[i]}-seed=${j}.log" 2>&1 &
+done