File size: 7,364 Bytes
7985e7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 | """
ProjectedDiffusionPolicy: IPOPT constraint projection at each DDPM step.
At each denoising step, takes a standard DDPM step then projects the result
onto the feasible set via IPOPT (when t <= projection_start threshold).
DDPM is batched on GPU; IPOPT projection is per-sample (CasADi limitation).
"""
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 ProjectedDiffusionPolicy(nn.Module):
"""DDPM sampling with per-step IPOPT projection.
Reverse-diffuses from t=1 to t=0, projecting the trajectory onto the
feasible set at each step where t <= projection_start.
"""
def __init__(
self,
model,
normalizer,
projector,
action_dim,
state_dim,
horizon,
n_sampling_steps=10,
projection_start=0.5,
n_projection_workers=1,
device="cuda:0",
):
super().__init__()
self.model = model
self.normalizer = normalizer
self.projector = projector
self.action_dim = action_dim
self.state_dim = state_dim
self.horizon = horizon
self.n_sampling_steps = n_sampling_steps
self.projection_start = projection_start
self.n_projection_workers = n_projection_workers
self.device = device
self.diffusion = CosineScheduleDiffusion(model, action_dim)
def __call__(self, conditions, batch_size=1):
"""Sample trajectories with projection.
Same interface as UnconditionalPolicy.__call__.
"""
conditions = apply_dict(
self.normalizer.normalize, conditions, "observations"
)
t_start = time.time()
transition_dim = self.action_dim + self.state_dim
conditions_torch = to_torch(conditions, device=self.device)
# start from pure noise
x = torch.randn(
batch_size, self.horizon, transition_dim, device=self.device
)
x = apply_conditioning(x, conditions_torch, self.action_dim)
# per-element IPOPT statistics accumulated over the denoise loop
self._ipopt_solves = 0
self._ipopt_iters = np.zeros(batch_size)
self._ipopt_failures = np.zeros(batch_size)
# time grid: t=1 (noisy) to t=0 (clean)
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)
# --- standard DDPM step ---
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, conditions_torch, self.action_dim)
if save_chain:
x0_estimates.append(x0_pred.detach().clone())
# DDPM posterior: q(x_{t_next} | x_{t_now}, x_0)
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, conditions_torch, self.action_dim)
# --- IPOPT projection ---
# projection_start is fraction of steps completed before projection begins
# 0.0 = all steps, 0.5 = second half, 1.0 = never
frac_done = i / self.n_sampling_steps
if frac_done >= self.projection_start:
x = self._project_batch(x, conditions_torch)
x = x.detach()
if save_chain:
denoising_chain.append(x.clone())
# unnormalize
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 self._ipopt_solves > 0:
info["ipopt_solves"] = np.full(batch_size, float(self._ipopt_solves))
info["ipopt_iters"] = self._ipopt_iters
info["ipopt_failures"] = self._ipopt_failures
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 _project_batch(self, x, conditions_torch):
"""Project each sample in the batch via IPOPT.
Args:
x: (B, H, transition_dim) normalized trajectory tensor
conditions_torch: dict with conditioned observations
Returns:
x_proj: (B, H, transition_dim) projected trajectory tensor
"""
B = x.shape[0]
x_np = to_np(x) # (B, H, transition_dim)
# convert full trajectories to DOF vectors
dofs = np.stack([self.projector.full_trajectory_to_dof(x_np[b]) for b in range(B)])
s0s = x_np[:, 0, self.action_dim:] # (B, state_dim)
# solve in parallel across workers
dofs_proj, statuses, iters = self.projector.project_batch(dofs, s0s, n_workers=self.n_projection_workers)
self._ipopt_solves += 1
self._ipopt_iters += np.asarray(iters, dtype=float)
self._ipopt_failures += np.asarray([s == "failed" for s in statuses], dtype=float)
# convert back to full trajectories
for b in range(B):
x_np[b] = self.projector.dof_to_full_trajectory(dofs_proj[b], s0s[b])
x_proj = torch.tensor(x_np, dtype=x.dtype, device=x.device)
# re-apply conditioning to ensure initial state is exact
x_proj = apply_conditioning(x_proj, conditions_torch, self.action_dim)
return x_proj
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)
|