File size: 3,660 Bytes
64b9802 | 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 | """GT motion-proxy ROI and GT-referenced dynamics. No actor/search sees future GT."""
import torch
import torch.nn.functional as F
from .rewards import CodecWorker
from .batch import TensorBatch
@torch.no_grad()
def motion_regions(truth,previous):
"""Shared across candidates: GT-derived, never shrinks after predicted disappearance."""
prev=torch.cat([previous[None],truth[:-1]],0)
motion=(truth-prev).abs().mean(1,keepdim=True)
thresholds=motion.flatten(1).quantile(.90,dim=1).clamp_min(.02)
masks=(motion>thresholds[:,None,None,None]).float()
masks=F.max_pool2d(masks,11,stride=1,padding=5)
fallback=masks.mean((1,2,3))<.01
masks[fallback]=1
crops=[];height,width=truth.shape[-2:]
for mask in masks:
y,x=torch.where(mask[0]>0)
y0,y1=int(y.min()),int(y.max())+1;x0,x1=int(x.min()),int(x.max())+1
cy=(y0+y1)//2;cx=(x0+x1)//2
h=max(64,y1-y0);w=max(64,x1-x0)
y0=max(0,min(cy-h//2,height-h));x0=max(0,min(cx-w//2,width-w))
crops.append((y0,min(height,y0+h),x0,min(width,x0+w)))
return masks,crops,fallback
class MotionReward:
def __init__(self,codec,perceptual,scales=None,alpha=.5,beta=.25):
self.worker=CodecWorker(codec,perceptual);self.perceptual=perceptual
self.scales=scales or dict(global_scale=1.,roi_scale=1.,delta_scale=1.)
self.alpha=alpha;self.beta=beta
@torch.no_grad()
def components(self,pred,pixels):
truth=pixels[0,2:];previous=pixels[0,1]
n,h=pred.shape[:2];real=truth[None].expand(n,-1,-1,-1,-1)
mae=(pred-real).abs().mean((2,3,4))
percept=self.worker.perceptual_loss(TensorBatch(dict(real=real,pred=pred))).batch['perceptual_loss']
masks,crops,fallback=motion_regions(truth,previous)
roi_mae=((pred-real).abs()*masks).sum((2,3,4))/(3*masks.sum((1,2,3)))[None]
roi_percept=torch.empty((n,h),device=pred.device)
for t,(y0,y1,x0,x1) in enumerate(crops):
r=truth[t:t+1,:,y0:y1,x0:x1]
for start in range(0,n,2):
pp=pred[start:start+2,t,:,y0:y1,x0:x1]
roi_percept[start:start+len(pp),t]=self.perceptual(pp*2-1,r.expand(len(pp),-1,-1,-1)*2-1).flatten()
pred_prev=torch.cat([previous[None,None].expand(n,1,-1,-1,-1),pred[:,:-1]],1)
true_prev=torch.cat([previous[None],truth[:-1]],0)
delta=((pred-pred_prev)-(truth-true_prev)[None]).abs().mean((2,3,4))
return dict(mae=mae,lpips=percept,global_loss=mae+percept,
roi_loss=roi_mae+roi_percept,roi_mae=roi_mae,roi_lpips=roi_percept,
delta_loss=delta,roi_fallback=fallback,roi_area=masks.mean((1,2,3)))
def combine(self,c):
s=self.scales
return -(c['global_loss']+self.alpha*s['global_scale']/s['roi_scale']*c['roi_loss']+
self.beta*s['global_scale']/s['delta_scale']*c['delta_loss']).mean(-1)
@torch.no_grad()
def decode(self,frames,context):
codes=torch.stack([f['response'] for f in frames],1)
assert codes.min()>=0 and codes.max()<4375
return self.worker.detokenize(TensorBatch(dict(tokens=codes,ctx_tokens=context.expand(len(codes),-1,-1))),None).batch['pixels'][:,1:]
@torch.no_grad()
def __call__(self,frames,context,actions,pixels):
c=self.components(self.decode(frames,context),pixels)
reward=self.combine(c)
assert torch.isfinite(reward).all()
metrics={'critic/recon_loss/mean':float(c['mae'].mean()),'critic/perceptual_loss/mean':float(c['lpips'].mean())}
metrics.update({k:float(v.float().mean()) for k,v in c.items()})
return reward,metrics
|