| """Causal ReCAP prompts shared by rollout and every policy log-probability pass.""" |
| import torch |
|
|
| def retained(target,recent=6): |
| assert target>=2 and recent>=1 |
| return [1]+list(range(max(2,target-recent),target)) |
|
|
| def make_prefix(context,history,actions,target,recent=6): |
| kept=retained(target,recent) |
| assert max(kept)==target-1 and len(kept)==len(set(kept)) |
| blocks=[context] |
| for t in kept:blocks += [history[t],actions[t].expand(context.shape[0],-1)] |
| return torch.cat(blocks,dim=1),kept |
|
|
| @torch.no_grad() |
| def rollout(model,context,anchor,actions,horizon,group,generator,visual=4375,width=80,temperature=1.0): |
| model.eval() |
| context=context.expand(group,-1);history={1:anchor.expand(group,-1)} |
| frames=[];traces=[] |
| for target in range(2,horizon+2): |
| prefix,kept=make_prefix(context,history,actions,target) |
| pending=prefix;kv=None;codes=[];lps=[] |
| for token in range(width): |
| out=model(input_ids=pending,past_key_values=kv,use_cache=True) |
| kv=out.past_key_values |
| logp=torch.log_softmax(out.logits[:,-1,:visual].float()/temperature,-1) |
| pending=torch.multinomial(logp.exp(),1,generator=generator) |
| codes.append(pending);lps.append(logp.gather(1,pending).squeeze(1)) |
| response=torch.cat(codes,1) |
| frames.append({'prefix':prefix.detach(),'response':response.detach(),'sample_logp':torch.stack(lps,1).detach(),'target':target}) |
| history[target]=response |
| traces.append({'target':target,'retained':kept,'prefix_length':prefix.shape[1],'positions_reset':True,'kv_reset':True,'future_observations':False}) |
| return frames,traces |
|
|
| def frame_logp(model,prefix,response,visual=4375,temperature=1.0): |
| |
| ids=torch.cat([prefix,response[:,:-1]],1) |
| out=model(input_ids=ids,use_cache=False) |
| p=prefix.shape[1] |
| logits=out.logits[:,p-1:p+response.shape[1]-1,:visual].float()/temperature |
| return torch.log_softmax(logits,-1).gather(-1,response.unsqueeze(-1)).squeeze(-1) |
|
|
