File size: 5,225 Bytes
0c3a82b | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """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)
|