| """Independent full chunks from one model, using immutable shared observations.""" |
| import argparse |
| import json |
| import logging |
| import time |
| import jax |
| import numpy as np |
| from serve_checkpoint import load_policy |
| from common import ROOT, EVAL, local_only, conditions, sha, write_json, save_npz, progress |
|
|
| def main(variant): |
| local_only() |
| manifest = json.loads((ROOT/'observations/manifest.json').read_text()) |
| for ident in conditions(): |
| assert sha(ROOT/'observations'/f'{ident}.npz') == manifest[ident]['sha256'] |
| start=time.monotonic() |
| wrapper=load_policy(EVAL.parent/'b1k_sft'/variant) |
| base=wrapper.base_policy |
| tap={}; transform=base._output_transform |
| def save_raw(outputs): |
| tap['raw']=np.asarray(outputs['actions'],dtype=np.float32).copy() |
| return transform(outputs) |
| base._output_transform=save_raw |
| outdir=ROOT/'samples'/variant; outdir.mkdir(parents=True,exist_ok=True) |
| norm=EVAL.parent/'b1k_sft'/variant/'checkpoint-14999/assets/IliaLarchenko/behavior_224_rgb/norm_stats.json' |
| norm_sha=sha(norm) |
| elapsed=[] |
| for ci, ident in enumerate(conditions()): |
| path=ROOT/'observations'/f'{ident}.npz' |
| meta=manifest[ident] |
| with np.load(path) as data: |
| obs={k:data[k] for k in data.files} |
| wrapper.current_stage=meta['stage'] |
| batch=wrapper.prepare_batch_for_pi_behavior(wrapper.process_obs(obs)) |
| assert int(batch['subtask_state']) == meta['stage'] |
| def draw(seed): |
| base._rng=jax.random.key(int(seed)) |
| t=time.monotonic() |
| result=base.infer(batch, initial_actions=None) |
| a=np.asarray(result['actions'],dtype=np.float32) |
| raw=tap['raw'].copy() |
| assert a.shape==(30,23) and raw.shape==(30,32) |
| assert np.isfinite(a).all() and np.isfinite(raw).all() |
| return a,raw,time.monotonic()-t,np.asarray(result['subtask_logits']) |
| if ci==0: |
| a,r,compile_s,logits=draw(90000) |
| b,s,_,_=draw(90000) |
| assert np.array_equal(a,b) and np.array_equal(r,s),'Same-seed repeatability failed' |
| check=[draw(90001+i) for i in range(32)] |
| preflight=dict(same_seed_exact=True,different_seed_physical_maxdiff=float(np.max(np.abs(check[0][0]-check[1][0]))), |
| different_seed_raw_maxdiff=float(np.max(np.abs(check[0][1]-check[1][1]))), |
| median_seconds=float(np.median([x[2] for x in check])),first_inference_seconds=compile_s, |
| load_and_preflight_seconds=time.monotonic()-start, normalizer_sha256=norm_sha, |
| initial_actions=None,num_flow_steps=20,wrapper_act_used=False,stage_fixed=True) |
| write_json(outdir/'preflight.json',preflight) |
| print('PREFLIGHT',json.dumps(preflight),flush=True) |
| final=outdir/f'{ident}.npz'; partial=outdir/f'{ident}.partial.npz' |
| actions=[]; raw_actions=[]; durations=[] |
| existing=final if final.exists() else partial |
| if existing.exists(): |
| with np.load(existing) as old: |
| assert str(old['observation_sha256'])==meta['sha256'] and str(old['normalizer_sha256'])==norm_sha |
| n=len(old['seeds']); assert np.array_equal(old['seeds'],np.arange(100000,100000+n)) |
| actions=list(old['actions']);raw_actions=list(old['raw_actions']);durations=list(old['seconds']) |
| if final.exists(): |
| assert n==512 |
| write_json(final.with_suffix('.json'),dict(variant=variant,condition=ident,samples=512, |
| sha256=sha(final),observation_sha256=meta['sha256'],normalizer_sha256=norm_sha, |
| num_flow_steps=20,initial_actions=None,physical_shape=[512,30,23],raw_shape=[512,30,32])) |
| continue |
| for j in range(len(actions),512): |
| a,r,d,logits=draw(100000+j) |
| actions.append(a);raw_actions.append(r);durations.append(d);elapsed.append(d) |
| if (j+1)%32==0: |
| save_npz(partial,actions=np.stack(actions),raw_actions=np.stack(raw_actions), |
| seeds=np.arange(100000,100001+j),seconds=np.array(durations), |
| observation_sha256=np.array(meta['sha256']),normalizer_sha256=np.array(norm_sha), |
| stage=np.array(meta['stage']),subtask_logits=logits) |
| progress('sampling_progress.json',variant=variant,condition=ident,condition_index=ci, |
| in_condition=j+1,model_samples=ci*512+j+1,model_target=10240, |
| median_seconds=float(np.median(elapsed[-128:]))) |
| partial.replace(final) |
| write_json(final.with_suffix('.json'),dict(variant=variant,condition=ident,samples=512, |
| sha256=sha(final),observation_sha256=meta['sha256'],normalizer_sha256=norm_sha, |
| num_flow_steps=20,initial_actions=None,physical_shape=[512,30,23],raw_shape=[512,30,32])) |
| print('SAMPLED',variant,ident,512,flush=True) |
| progress(f'logs/{variant}_sampling_complete.json',state='complete',samples=10240) |
|
|
| if __name__=='__main__': |
| logging.basicConfig(level=logging.INFO) |
| p=argparse.ArgumentParser();p.add_argument('variant');main(p.parse_args().variant) |
|
|
