| """Record policy inputs before wrapper stage updates; serve ordinary reference control.""" |
| import argparse |
| import asyncio |
| import http |
| import logging |
| import time |
| import jax |
| import numpy as np |
| from websockets.asyncio.server import serve |
| from openpi_client import msgpack_numpy |
| from serve_checkpoint import load_policy |
| from common import ROOT, EVAL, PORT, local_only, progress, save_npz, write_json, sha |
|
|
| async def main(variant): |
| local_only() |
| policy = load_policy(EVAL.parent / 'b1k_sft' / variant) |
| dummy = {'robot_r1::proprio': np.zeros(61, np.float32), 'task_id': np.array([0])} |
| for name in ['zed_link', 'left_realsense_link', 'right_realsense_link']: |
| dummy[f'robot_r1::robot_r1:{name}:Camera:0::rgb'] = np.zeros((224, 224, 3), np.uint8) |
| for _ in range(27): |
| assert np.isfinite(policy.act(dummy).numpy()).all() |
| policy.reset() |
| policy.base_policy._rng = jax.random.key(0) |
| last = None |
|
|
| def snapshot(obs, meta, stage): |
| ident = meta['condition'] |
| path = ROOT / 'observations' / f'{ident}.npz' |
| if path.exists(): |
| return |
| arrays = {k: np.asarray(v) for k, v in obs.items() if k == 'robot_r1::proprio' or k.endswith('::rgb') or k == 'task_id'} |
| assert arrays['robot_r1::proprio'].shape == (61,) |
| save_npz(path, **arrays) |
| write_json(path.with_suffix('.json'), dict(meta, stage=int(stage), source_model=variant, |
| seed=0, task=0, sha256=sha(path), created=time.time(), |
| step_convention='number of completed environment steps; input before next control action', |
| shapes={k:list(v.shape) for k,v in arrays.items()})) |
| print('SAVED', ident, 'stage', stage, flush=True) |
|
|
| async def handler(ws): |
| nonlocal last |
| await ws.send(msgpack_numpy.packb({'task_id': 0, 'variant': variant})) |
| async for message in ws: |
| obs = msgpack_numpy.unpackb(message) |
| if 'reset' in obs: |
| policy.reset() |
| policy.base_policy._rng = jax.random.key(0) |
| last = None |
| continue |
| meta = obs.pop('_capture', {}) |
| if meta.get('fallback'): |
| assert last is not None |
| old_obs, old_stage, old_meta = last |
| snapshot(old_obs, dict(meta, actual_step=old_meta['step'], reason='reference episode terminated before requested step'), old_stage) |
| await ws.send(msgpack_numpy.packb({'action': np.zeros(23, np.float32)})) |
| continue |
| if meta.get('condition'): |
| snapshot(obs, dict(meta, actual_step=meta['step']), policy.current_stage) |
| last = (obs, policy.current_stage, meta) |
| action = policy.act(obs).numpy() |
| assert action.shape == (23,) and np.isfinite(action).all() |
| if meta.get('step', 0) % 100 == 0: |
| progress('capture_progress.json', variant=variant, instance=meta.get('instance'), step=meta.get('step'), stage=policy.current_stage) |
| await ws.send(msgpack_numpy.packb({'action': action})) |
|
|
| def health(connection, request): |
| if request.path == '/healthz': |
| return connection.respond(http.HTTPStatus.OK, 'OK\n') |
| async with serve(handler, '127.0.0.1', PORT, compression=None, max_size=None, |
| ping_interval=None, process_request=health) as server: |
| progress(f'logs/{variant}_ready.json', state='ready') |
| await server.serve_forever() |
|
|
| if __name__ == '__main__': |
| logging.basicConfig(level=logging.INFO) |
| p = argparse.ArgumentParser(); p.add_argument('variant') |
| asyncio.run(main(p.parse_args().variant)) |
|
|
