|
|
import torch as t |
|
|
|
|
|
@t.no_grad() |
|
|
def sample(v, z, actions, num_steps=10, cfg=0, negative_actions=None, cache=None): |
|
|
return sample_with_grad(v, z, actions, num_steps, cfg, negative_actions, cache=cache) |
|
|
|
|
|
def sample_with_grad(v, z, actions, num_steps=10, cfg=0, negative_actions=None, cache=None): |
|
|
device = v.device |
|
|
ts = 1 - t.linspace(0, 1, num_steps+1, device=device) |
|
|
ts = 3*ts/(2*ts + 1) |
|
|
z_prev = z.clone() |
|
|
z_prev = z_prev.to(device) |
|
|
for i in range(len(ts)-1): |
|
|
t_cond = ts[i].repeat(z_prev.shape[0], 1) |
|
|
cached_k = None |
|
|
cached_v = None |
|
|
if cache is not None: |
|
|
cached_k, cached_v = cache.get() |
|
|
|
|
|
v_pred, k_new, v_new = v(z_prev.to(device), actions.to(device), t_cond.to(device), cached_k=cached_k, cached_v=cached_v) |
|
|
|
|
|
if i == len(ts)-2 and cache is not None: |
|
|
cache.extend(k_new, v_new) |
|
|
|
|
|
if cfg > 0: |
|
|
if cache is not None: |
|
|
raise NotImplementedError("this is not implemented yet") |
|
|
if negative_actions is not None: |
|
|
v_neg, _, _ = v(z_prev.to(device), negative_actions.to(device), t_cond.to(device)) |
|
|
else: |
|
|
v_neg, _, _ = v(z_prev.to(device), t.zeros_like(actions, dtype=t.long, device=device), t_cond.to(device)) |
|
|
v_pred = v_neg + cfg * (v_pred - v_neg) |
|
|
z_prev = z_prev + (ts[i] - ts[i+1])*v_pred |
|
|
return z_prev |
