| """Four disjoint beam groups, with unmodified policy scores for the RL surrogate. |
| |
| Search penalties only affect proposal selection. Deterministic proposals do not |
| have on-policy sampling probabilities; the beam RL objective remains biased. |
| """ |
| import torch |
| from .memory import make_prefix |
|
|
| @torch.no_grad() |
| def rollout_group_beam(model, context, anchor, actions, horizon, penalty, |
| groups=4, beams=2, visual=4375, width=80): |
| model.eval() |
| device=context.device; count=groups*beams |
| history={1:anchor.expand(count,-1)}; frames=[]; traces=[] |
| scores=torch.full((count,),float('-inf'),device=device,dtype=torch.float64) |
| scores[::beams]=0 |
| raw=torch.zeros(count,device=device,dtype=torch.float64) |
| for target in range(2,horizon+2): |
| original,kept=make_prefix(context.expand(count,-1),history,actions,target) |
| ancestors=torch.arange(count,device=device) |
| codes=torch.empty((count,0),device=device,dtype=torch.long) |
| logs=torch.empty((count,0),device=device) |
| pending=original; cache=None |
| for token in range(width): |
| out=model(input_ids=pending,past_key_values=cache,use_cache=True) |
| lp=out.logits[:,-1,:visual].float().log_softmax(-1) |
| parents=[]; choices=[]; next_scores=[] |
| for g in range(groups): |
| start=g*beams |
| if g: |
| frequency=torch.bincount(torch.cat(choices),minlength=visual).double()/(g*beams) |
| else: |
| frequency=torch.zeros(visual,device=device,dtype=torch.float64) |
| totals=scores[start:start+beams,None]+lp[start:start+beams].double()-penalty*frequency |
| score,flat=totals.flatten().topk(beams,sorted=True) |
| parents.append(start+flat//visual);choices.append(flat%visual);next_scores.append(score) |
| parent=torch.cat(parents);chosen=torch.cat(choices) |
| assert torch.equal(parent//beams,torch.arange(count,device=device)//beams) |
| raw=raw[parent]+lp[parent,chosen].double() |
| scores=torch.cat(next_scores) |
| codes=torch.cat([codes[parent],chosen[:,None]],1) |
| logs=torch.cat([logs[parent],lp[parent,chosen,None]],1) |
| ancestors=ancestors[parent] |
| cache=out.past_key_values;cache.reorder_cache(parent) |
| pending=chosen[:,None] |
| for f in frames: |
| for key in ['prefix','response','model_logp']:f[key]=f[key][ancestors] |
| history={t:x[ancestors] for t,x in history.items()} |
| history[target]=codes |
| frames.append(dict(target=target,prefix=original[ancestors],response=codes,model_logp=logs)) |
| traces.append(dict(target=target,retained=kept,prefix_length=original.shape[1], |
| frame_parent_indices=ancestors.tolist(),positions_reset=True,kv_reset=True, |
| future_observations=False,random_sampling=False)) |
| selected=torch.arange(0,count,beams,device=device) |
| frames=[{k:(v[selected] if torch.is_tensor(v) else v) for k,v in f.items()} for f in frames] |
| history={t:x[selected] for t,x in history.items()} |
| for f in frames: |
| rebuilt,_=make_prefix(context.expand(groups,-1),history,actions,f['target']) |
| assert torch.equal(rebuilt,f['prefix']) |
| rebuilt=sum(f['model_logp'].double().sum(-1) for f in frames) |
| assert torch.allclose(rebuilt,raw[selected],atol=1e-7,rtol=0) |
| codes=torch.stack([f['response'] for f in frames],1) |
| pair=[float((codes[i]!=codes[j]).float().mean()) for i in range(groups) for j in range(i+1,groups)] |
| perframe=(codes!=codes[:1]).any(0).sum(-1).tolist() |
| return frames,dict(strategy='grouped_diverse_full_video_beam',groups=groups,beams_per_group=beams, |
| active_paths=count,returned_trajectories=groups,penalty=penalty, |
| unique_trajectories=len({tuple(c.flatten().tolist()) for c in codes}), |
| mean_pairwise_token_difference_fraction=sum(pair)/len(pair), |
| pairwise_token_difference_fraction=pair,different_token_positions=sum(perframe), |
| different_positions_per_frame=perframe,model_log_scores=raw[selected].tolist(), |
| search_scores=scores[selected].tolist(),score_is_behavior_logprob=False, |
| on_policy_grpo=False,frames=traces) |
|
|
