| import math |
| import time |
| from collections import namedtuple |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| from model import CosineScheduleDiffusion, apply_conditioning |
| from utils.arrays import to_np, to_torch, apply_dict |
|
|
| Trajectories = namedtuple("Trajectories", "actions observations values") |
|
|
|
|
| class UnconditionalPolicy(nn.Module): |
| """DDPM sampling policy — raw model outputs, no guidance. |
| |
| Reverse-diffuses from t=1 (noise) to t=0 (data) using the cosine schedule |
| and DDPM posterior steps. The last step is always deterministic. |
| |
| Supports batch_size > 1 for parallel episode evaluation. |
| """ |
|
|
| def __init__( |
| self, |
| model, |
| normalizer, |
| action_dim, |
| state_dim, |
| horizon, |
| n_sampling_steps=10, |
| device="cuda:0", |
| ): |
| super().__init__() |
| self.model = model |
| self.normalizer = normalizer |
| self.action_dim = action_dim |
| self.state_dim = state_dim |
| self.horizon = horizon |
| self.n_sampling_steps = n_sampling_steps |
| self.device = device |
|
|
| self.diffusion = CosineScheduleDiffusion(model, action_dim) |
|
|
| def __call__(self, conditions, batch_size=1): |
| """Sample trajectories via DDPM reverse process. |
| |
| Args: |
| conditions: {0: obs} where obs is (state_dim,) or (batch_size, state_dim) |
| batch_size: number of trajectories to sample in parallel |
| |
| Returns: |
| actions: (batch_size, action_dim) — first action per trajectory |
| trajectories: Trajectories namedtuple with (batch_size, horizon, ...) arrays |
| denoising_chain: unnormalized chain (only for batch_size=1, else None) |
| x0_estimates: unnormalized estimates (only for batch_size=1, else None) |
| info: dict with computation_time |
| """ |
| |
| conditions = apply_dict( |
| self.normalizer.normalize, conditions, "observations" |
| ) |
|
|
| t_start = time.time() |
|
|
| transition_dim = self.action_dim + self.state_dim |
|
|
| |
| x = torch.randn( |
| batch_size, self.horizon, transition_dim, device=self.device |
| ) |
| x = apply_conditioning(x, to_torch(conditions, device=self.device), self.action_dim) |
|
|
| |
| timesteps = torch.linspace(1.0, 0.0, self.n_sampling_steps + 1, device=self.device) |
|
|
| save_chain = (batch_size == 1) |
| if save_chain: |
| denoising_chain = [x.clone()] |
| x0_estimates = [] |
|
|
| for i in range(self.n_sampling_steps): |
| t_now = timesteps[i] |
| t_next = timesteps[i + 1] |
| is_last_step = (i == self.n_sampling_steps - 1) |
|
|
| abar_now = self.diffusion.alpha_bar(t_now) |
| abar_next = self.diffusion.alpha_bar(t_next) |
|
|
| |
| t_batch = t_now.expand(batch_size) |
| x0_pred = self.model(x, t_batch) |
|
|
| |
| x0_pred = apply_conditioning(x0_pred, to_torch(conditions, device=self.device), self.action_dim) |
|
|
| if save_chain: |
| x0_estimates.append(x0_pred.detach().clone()) |
|
|
| |
| beta = 1.0 - abar_now / abar_next |
|
|
| coeff_x0 = math.sqrt(abar_next) * beta / (1.0 - abar_now) |
| coeff_xt = math.sqrt(1.0 - beta) * (1.0 - abar_next) / (1.0 - abar_now) |
| posterior_mean = coeff_x0 * x0_pred + coeff_xt * x |
|
|
| if is_last_step: |
| x = posterior_mean |
| else: |
| posterior_var = max((1.0 - abar_next) / (1.0 - abar_now) * beta, 0.0) |
| noise = torch.randn_like(x) |
| x = posterior_mean + math.sqrt(posterior_var) * noise |
|
|
| |
| x = apply_conditioning(x, to_torch(conditions, device=self.device), self.action_dim) |
| x = x.detach() |
|
|
| if save_chain: |
| denoising_chain.append(x.clone()) |
|
|
| |
| normed_actions = x[:, :, :self.action_dim] |
| actions = self.normalizer.unnormalize(to_np(normed_actions), "actions") |
|
|
| normed_observations = x[:, :, self.action_dim:] |
| observations = self.normalizer.unnormalize(to_np(normed_observations), "observations") |
|
|
| trajectories = Trajectories(actions, observations, np.zeros((batch_size, 1))) |
|
|
| t_end = time.time() |
| info = {"computation_time": t_end - t_start} |
|
|
| if save_chain: |
| chain = torch.stack(denoising_chain, dim=1) |
| chain = self._unnormalize_chain(chain) |
| x0_est = torch.stack(x0_estimates, dim=1) |
| x0_est = self._unnormalize_chain(x0_est) |
| return actions[:, 0], trajectories, chain, x0_est, info |
| else: |
| return actions[:, 0], trajectories, None, None, info |
|
|
| def _unnormalize_chain(self, chain): |
| """Unnormalize a chain of shape (batch, steps, horizon, transition_dim).""" |
| obs_chain = chain[..., self.action_dim:] |
| obs_unnorm = self.normalizer.unnormalize(to_np(obs_chain), "observations") |
| act_chain = chain[..., :self.action_dim] |
| act_unnorm = self.normalizer.unnormalize(to_np(act_chain), "actions") |
| return np.concatenate([act_unnorm, obs_unnorm], axis=-1) |
|
|
