File size: 7,096 Bytes
68eed61 | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """Bounded, isolated diagnostics; never changes or restarts the training job."""
import argparse
import dataclasses
import json
import os
from pathlib import Path
import subprocess
import time
import urllib.request
ROOT = Path(__file__).resolve().parents[1]
COSMOS = Path('/home/gpu4/tianyi/cosmos-framework')
PYTHON = '/home/gpu4/tianyi/cosmos-framework/.venv/bin/python'
SIM = '/home/gpu4/miniconda3/envs/simpler_env/bin/python'
def closed_loop(args):
# Run foreground under this bounded parent; always reap our own server.
with (args.output / 'server.log').open('w') as log:
server = subprocess.Popen([PYTHON, str(ROOT/'scripts/widowx_policy_server.py'),
'--delta', str(args.checkpoint), '--port', str(args.port)],
stdout=log, stderr=subprocess.STDOUT, cwd=COSMOS)
try:
deadline = time.monotonic() + 240
while True:
if server.poll() is not None:
raise RuntimeError(f'server exited {server.returncode}')
try:
with urllib.request.urlopen(f'http://127.0.0.1:{args.port}/healthz', timeout=2) as r:
if r.status == 200:
break
except OSError:
pass
if time.monotonic() >= deadline:
raise TimeoutError('server not ready')
time.sleep(2)
env = dict(os.environ, PYTHONPATH='/opt/benchmarks/SimplerEnv',
MS2_ASSET_DIR='/opt/benchmarks/SimplerEnv/ManiSkill2_real2sim/data',
VK_ICD_FILENAMES='/etc/vulkan/icd.d/nvidia_icd.json')
for convention in ('bridge_tool', 'legacy_tcp'):
destination = args.output / convention
destination.mkdir(exist_ok=True)
with (destination/'episodes.log').open('w') as elog:
subprocess.run([SIM, str(ROOT/'scripts/eval_widowx_simpler.py'),
'--port', str(args.port), '--checkpoint', str(args.checkpoint),
'--output', str(destination/'results.json'), '--episodes', '10',
'--proprio-convention', convention,
'--video-dir', str(destination/'videos'),
'--trace-dir', str(destination/'traces')],
env=env, stdout=elog, stderr=subprocess.STDOUT, check=True, timeout=600,
cwd=COSMOS)
report=json.loads((destination/'results.json').read_text())
print(convention, report['success_rate'], flush=True)
finally:
server.terminate()
try:
server.wait(timeout=20)
except subprocess.TimeoutExpired:
server.kill()
server.wait()
def offline(args):
import numpy as np
import torch
from scripts.widowx_policy_server import WidowXService
from cosmos_framework.scripts.action_policy_server_robocasa365_zeva import RobolabServerArgs
from cosmos_framework.data.generator.action.datasets.widowx_bridge_v3_dataset import WidowXBridgeV3Dataset, normalize_actions
torch.set_num_threads(2)
release=Path('/home/gpu4/tianyi/zeva-release/weights')
os.chdir(COSMOS)
service=WidowXService(RobolabServerArgs(
checkpoint_path=str(release/'stage2'), allow_dcp_checkpoint=True,
experiment='action_policy_robocasa365_atomic5_zeva',
experiment_overrides=['model.config.tokenizer.vae_path='+os.environ['WAN_VAE_PATH'],
'model.config.vlm_config.tokenizer.pretrained_model_name='+os.environ['QWEN_VLM_PATH']],
task_context_bank=release/'stage1/train_memory_effect_v3.pt',
cte_checkpoint=release/'stage1/zeva_cte.pt', static_task_context_checkpoint=release/'stage3/best.pt',
domain_name='bridge_orig_lerobot', image_height=256, image_width=256,
resolution='256', conditioning_fps=5, action_dim=7, use_state=False))
dataset=WidowXBridgeV3Dataset('/opt/zhangchenyu/datasets/bridge_orig_lerobot_smoke',
feature_cache=str(ROOT/'.artifacts/widowx/features'))
indices=[0, dataset.windows.index((306,0)), len(dataset)-1,
next(i for i,(_,s) in enumerate(dataset.windows) if s>=16)]
named=dict(service.model.named_parameters())
metrics=[]
for label, checkpoint in [('step2',ROOT/'.artifacts/widowx/checkpoints/step_000000002.pt'),
('trained',args.checkpoint)]:
payload=torch.load(checkpoint,map_location='cpu',weights_only=False)
with torch.no_grad():
for name,value in payload['model'].items():
p=named[name]; (p.to_local() if hasattr(p,'to_local') else p).copy_(value)
for i in indices:
sample=dataset[i]; ep,start=dataset.windows[i]
frames=dataset._frames[ep]; rows=dataset._rows[ep]
request={'prompt':sample['ai_caption'], 'observation/image':frames[start],
'observation/proprio':sample['proprio'].numpy(),
'cte_boundary_images':frames[:start+1:4],
'cte_transition_actions':normalize_actions(rows['action'][:start]).reshape(-1,4,7),
'inference_seed':0}
for steps,guidance in ((4,3.0),(16,1.0)) if label=='trained' else ((4,3.0),):
service.cfg=dataclasses.replace(service.cfg,num_steps=steps,guidance=guidance)
result=service.infer(request)
prediction=np.asarray(result['action'],np.float32)
target=sample['action'].numpy()
error=np.abs(prediction-target)
item=dict(checkpoint=label,iteration=payload['iteration'],episode=ep,start=start,
steps=steps,guidance=guidance,mae=float(error.mean()),
first4_mae=float(error[:4].mean()),per_dim_mae=error.mean(0).tolist(),
gripper_accuracy=float(np.mean((prediction[:,6]>0)==(target[:,6]>0))),
out_of_range_fraction=float(np.mean(np.abs(prediction)>1)))
metrics.append(item)
np.savez(args.output/f'{label}_ep{ep}_start{start}_steps{steps}.npz',
prediction=prediction,target=target,proprio=sample['proprio'].numpy())
print(json.dumps(item),flush=True)
(args.output/'metrics.json').write_text(json.dumps(metrics,indent=2))
if __name__=='__main__':
parser=argparse.ArgumentParser()
parser.add_argument('--mode',choices=('closed_loop','offline'),required=True)
parser.add_argument('--checkpoint',type=Path,required=True)
parser.add_argument('--output',type=Path,required=True)
parser.add_argument('--port',type=int,default=18766)
args=parser.parse_args(); args.output.mkdir(parents=True,exist_ok=True)
globals()[args.mode](args)
|