| """Evaluate Zeva deltas continuously, or run one explicitly selected checkpoint.""" |
| import argparse |
| import json |
| import os |
| import socket |
| from pathlib import Path |
| import subprocess |
| import time |
| import urllib.request |
|
|
| ROOT=Path(__file__).resolve().parents[1] |
| |
| |
| RUN=Path(os.environ.get('WIDOWX_RUN_DIR', str(ROOT/'.artifacts/widowx'))) |
| PYTHON='/home/gpu4/tianyi/cosmos-framework/.venv/bin/python' |
| SIM_PYTHON='/home/gpu4/miniconda3/envs/simpler_env/bin/python' |
| PORT=18765 |
| EVAL_GPU=os.environ.get('WIDOWX_EVAL_GPU','3') |
| RENDER_GPU=os.environ.get('WIDOWX_RENDER_GPU', EVAL_GPU) |
|
|
| def run_one(checkpoint, episodes=10, port=PORT, record_traces=False, max_steps=None, num_steps=4, seed_start=0, exec_horizon=4, guidance=3.0, memory_mode='on'): |
| |
| |
| with socket.socket() as probe: |
| probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) |
| probe.bind(('127.0.0.1', port)) |
| out=RUN/'eval'; out.mkdir(parents=True,exist_ok=True) |
| gpu_uuid=subprocess.check_output(['nvidia-smi','-i',EVAL_GPU,'--query-gpu=uuid', |
| '--format=csv,noheader'],text=True).strip() |
| if not gpu_uuid.startswith('GPU-') or '\n' in gpu_uuid: |
| raise ValueError('WIDOWX_EVAL_GPU must select exactly one physical GPU') |
| render_uuid=subprocess.check_output(['nvidia-smi','-i',RENDER_GPU,'--query-gpu=uuid', |
| '--format=csv,noheader'],text=True).strip() |
| if not render_uuid.startswith('GPU-') or '\n' in render_uuid: |
| raise ValueError('WIDOWX_RENDER_GPU must select exactly one physical GPU') |
| env=os.environ.copy() |
| env.update(CUDA_VISIBLE_DEVICES=gpu_uuid,PYTHONPATH=str(ROOT)+':/home/gpu4/tianyi/cosmos-framework', |
| WIDOWX_POLICY_OUTPUT_DIR=str(out/'server_runtime'), |
| QWEN_VLM_PATH='/home/gpu4/tianyi/models/Qwen3-VL-8B-Instruct', |
| WAN_VAE_PATH='/home/gpu4/tianyi/models/wan22_vae/Wan2.2_VAE.pth', |
| HF_HUB_OFFLINE='1',TOKENIZERS_PARALLELISM='false') |
| simenv=env.copy() |
| simenv.update(CUDA_VISIBLE_DEVICES=render_uuid,PYTHONPATH='/opt/benchmarks/SimplerEnv', |
| VK_ICD_FILENAMES='/etc/vulkan/icd.d/nvidia_icd.json', |
| MS2_ASSET_DIR='/opt/benchmarks/SimplerEnv/ManiSkill2_real2sim/data') |
| subprocess.run([SIM_PYTHON,str(ROOT/'scripts/check_sapien_gpu.py'), |
| '--expected-uuid',render_uuid],env=simenv,check=True) |
| with (out/(checkpoint.stem+'.server.log')).open('a') as log: |
| server=subprocess.Popen([PYTHON,str(ROOT/'scripts/widowx_policy_server.py'), |
| '--delta',str(checkpoint),'--port',str(port),'--num-steps',str(num_steps), |
| '--guidance',str(guidance),'--memory-mode',memory_mode], |
| env=env,cwd='/home/gpu4/tianyi/cosmos-framework',stdout=log,stderr=subprocess.STDOUT) |
| try: |
| deadline=time.monotonic()+600 |
| while True: |
| if server.poll() is not None: raise RuntimeError(f'policy server exited {server.returncode}') |
| try: |
| with urllib.request.urlopen(f'http://127.0.0.1:{port}/healthz',timeout=3) as r: |
| if r.status==200: break |
| except Exception: pass |
| if time.monotonic()>deadline: raise TimeoutError('policy server startup') |
| time.sleep(3) |
| trace_args = ['--trace-dir', str(out/'traces'/checkpoint.stem)] if record_traces else [] |
| limit_args = ['--max-steps', str(max_steps)] if max_steps is not None else [] |
| with (out/(checkpoint.stem+'.episodes.log')).open('a') as eval_log: |
| subprocess.run([SIM_PYTHON,str(ROOT/'scripts/eval_widowx_simpler.py'), |
| '--checkpoint',str(checkpoint),'--output',str(out/(checkpoint.stem+'.json')), |
| '--video-dir',str(out/'videos'/checkpoint.stem), |
| '--episodes',str(episodes),'--port',str(port),'--seed-start',str(seed_start), |
| '--exec-horizon',str(exec_horizon), |
| '--memory-mode',memory_mode, |
| *trace_args,*limit_args],env=simenv,check=True, |
| stdout=eval_log,stderr=subprocess.STDOUT) |
| report=json.loads((out/(checkpoint.stem+'.json')).read_text()) |
| if report.get('status') != 'complete' or len(report.get('episodes', [])) != episodes: |
| raise RuntimeError('evaluation did not produce the requested number of complete episodes') |
| report['inference_num_steps'] = num_steps |
| report['guidance'] = guidance |
| report['physical_gpu_uuid'] = gpu_uuid |
| report['render_gpu_uuid'] = render_uuid |
| report['memory_mode'] = memory_mode |
| report_path = out/(checkpoint.stem+'.json') |
| temporary = report_path.with_suffix('.tmp') |
| temporary.write_text(json.dumps(report, indent=2)) |
| temporary.replace(report_path) |
| print('EVAL_COMPLETE '+json.dumps(dict(checkpoint=str(checkpoint), |
| episodes=episodes,successes=sum(bool(x['success']) for x in report['episodes']), |
| success_rate=report['success_rate'])),flush=True) |
| finally: |
| server.terminate() |
| try: server.wait(timeout=30) |
| except subprocess.TimeoutExpired: server.kill(); server.wait() |
|
|
| def main(): |
| parser=argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--once', action='store_true', help='evaluate one checkpoint and exit; never starts training') |
| parser.add_argument('--checkpoint', type=Path, help='explicit checkpoint for --once') |
| parser.add_argument('--episodes', type=int, default=10) |
| parser.add_argument('--port', type=int, default=PORT) |
| parser.add_argument('--record-traces', action='store_true') |
| parser.add_argument('--num-steps', type=int, default=4, help='UniPC denoising steps per prediction') |
| parser.add_argument('--seed-start', type=int, default=0, help='first environment and inference seed') |
| parser.add_argument('--exec-horizon', type=int, default=4, choices=range(1,33)) |
| parser.add_argument('--guidance', type=float, default=3.0) |
| parser.add_argument('--memory-mode',choices=('on','bank_off'),default='on') |
| parser.add_argument('--max-steps', type=int, default=None, |
| help='maximum environment control steps per episode') |
| args=parser.parse_args() |
| if args.episodes < 1 or not 1 <= args.port <= 65535 or (args.max_steps is not None and args.max_steps < 1): |
| parser.error('episodes and max-steps must be positive; port must be in [1,65535]') |
| if args.num_steps < 1 or args.seed_start < 0: |
| parser.error('num-steps must be positive and seed-start must be nonnegative') |
| if args.once != (args.checkpoint is not None): |
| parser.error('--once and --checkpoint must be supplied together') |
| (RUN/'eval').mkdir(parents=True,exist_ok=True) |
| if args.once: |
| checkpoint=args.checkpoint.resolve() |
| if not checkpoint.is_file(): |
| parser.error(f'checkpoint does not exist: {checkpoint}') |
| if (RUN/'eval'/(checkpoint.stem+'.json')).exists(): |
| parser.error('a result already exists for this checkpoint; use a separate WIDOWX_RUN_DIR') |
| print('EVALUATING_ONCE '+json.dumps(dict(checkpoint=str(checkpoint), |
| episodes=args.episodes,gpu=EVAL_GPU,port=args.port, |
| inference_num_steps=args.num_steps,seed_start=args.seed_start,max_steps=args.max_steps, |
| memory_mode=args.memory_mode,render_gpu=RENDER_GPU)),flush=True) |
| try: |
| run_one(checkpoint, args.episodes, args.port, args.record_traces, args.max_steps, args.num_steps, args.seed_start, args.exec_horizon, args.guidance, args.memory_mode) |
| except Exception as e: |
| (RUN/'eval'/(checkpoint.stem+'.error.json')).write_text( |
| json.dumps(dict(status='failed',error=repr(e),time=time.time()))) |
| raise |
| return |
| while True: |
| |
| |
| for checkpoint in sorted((RUN/'checkpoints').glob('step_*.pt'), reverse=True): |
| if int(checkpoint.stem.split('_')[1]) % 2000: |
| continue |
| if (RUN/'eval'/(checkpoint.stem+'.json')).exists(): continue |
| try: |
| print('EVALUATING',checkpoint,flush=True) |
| run_one(checkpoint, args.episodes, args.port, args.record_traces, args.max_steps, args.num_steps, args.seed_start, args.exec_horizon, args.guidance, args.memory_mode) |
| except Exception as e: |
| failure=RUN/'eval'/(checkpoint.stem+'.error.json') |
| failure.write_text(json.dumps(dict(status='failed',error=repr(e),time=time.time()))) |
| print('EVAL_FAILED',repr(e),flush=True) |
| time.sleep(30) |
| time.sleep(10) |
|
|
| if __name__=='__main__': main() |
|
|
