File size: 6,466 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 | """Cache frozen release CTE/retrieval features using only each window's past."""
import os
import argparse
from pathlib import Path
import numpy as np
import torch
from cosmos_framework.data.generator.action.datasets.widowx_bridge_v3_dataset import WidowXBridgeV3Dataset, normalize_actions
from cosmos_framework.scripts.action_policy_server_robocasa365_zeva import RobolabServerArgs, RobolabPolicyService
from scripts.widowx_contract import ensure_manifest
@torch.inference_mode()
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--start-episode",
type=int,
default=None,
help="inclusive episode index to process (default: first episode)",
)
parser.add_argument(
"--end-episode",
type=int,
default=None,
help="exclusive episode index to process (default: after last episode)",
)
parser.add_argument('--shard-index', type=int, default=0)
parser.add_argument('--num-shards', type=int, default=1)
parser.add_argument('--validate-only', action='store_true')
args = parser.parse_args()
if not 0 <= args.shard_index < args.num_shards:
parser.error('require 0 <= shard-index < num-shards')
release = Path('/home/gpu4/tianyi/zeva-release/weights')
out = Path(os.environ['WIDOWX_FEATURE_CACHE']); out.mkdir(parents=True, exist_ok=True)
cache_fingerprint = ensure_manifest(os.environ['WIDOWX_DATASET'], out)
ds = WidowXBridgeV3Dataset(os.environ['WIDOWX_DATASET'], window_mode='full')
lengths = {e['episode_index']: int(e['length']) for e in ds.episodes}
episodes = sorted(lengths)[args.shard_index::args.num_shards]
if args.start_episode is not None:
episodes = [e for e in episodes if e >= args.start_episode]
if args.end_episode is not None:
episodes = [e for e in episodes if e < args.end_episode]
if not episodes:
raise SystemExit('no episodes selected')
pending = []
for ep in episodes:
path = out / f'episode_{ep:06d}.pt'
if not path.exists():
pending.append(ep)
continue
payload = torch.load(path, weights_only=True, map_location='cpu')
starts = list(range(0, lengths[ep], 4))
if payload.get('fingerprint') != cache_fingerprint or payload.get('episode_index') != ep:
raise ValueError(f'Cache provenance mismatch: {path}')
if payload['starts'].tolist() != starts:
raise ValueError(f'Incomplete episode cache: {path}')
for key, tail in [('behavior_global',(256,)), ('behavior_phase',(128,)),
('behavior_effect',(4,128)), ('behavior_effect_valid',(4,))]:
if tuple(payload[key].shape) != (len(starts), *tail) or not torch.isfinite(payload[key]).all():
raise ValueError(f'Invalid feature tensor {key}: {path}')
if args.validate_only:
if pending:
raise SystemExit(f'Feature cache incomplete: {len(pending)} missing episodes')
print(f'FEATURE_CACHE_VALID episodes={len(episodes)} fingerprint={cache_fingerprint}', flush=True)
return
episodes = pending
if not episodes:
print('FEATURE_CACHE_ALREADY_COMPLETE', flush=True)
return
service = RobolabPolicyService(RobolabServerArgs(
checkpoint_path=os.environ.get('BASE_CHECKPOINT_PATH', str(release/'stage2')), allow_dcp_checkpoint=True,
output_dir=out.parent/'feature_runtime'/f'shard_{args.shard_index}',
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,
conditioning_fps=5, action_dim=7, use_state=False))
for count, ep in enumerate(episodes):
path = out / f'episode_{ep:06d}.pt'
ds._load(ep)
rgb, rows = ds._frames[ep], ds._rows[ep]
starts = boundary = list(range(0, len(rgb), 4))
latents = torch.stack([service._encode_cte_frame(rgb[i]) for i in boundary]).unsqueeze(0)
raw = normalize_actions(rows['action'])
transitions = torch.from_numpy(raw[:boundary[-1]].reshape(-1,4,7)).to(latents.device).unsqueeze(0)
valid = torch.ones(latents.shape[:2], dtype=torch.bool, device=latents.device)
# The released CTE is causal; run once per prefix to also enforce causality
# if a future implementation changes its internal attention masking.
features = {k: [] for k in ('behavior_global','behavior_phase','behavior_effect','behavior_effect_valid')}
instruction = ds.tasks[int(rows['task_index'][0])]
global_value = service._task_context_from_initial_observation(rgb[0], instruction)[0].float().cpu()
for j in range(len(boundary)):
a = transitions[:, :j]
with torch.autocast('cuda', dtype=torch.bfloat16):
enc = service._cte(latents[:, :j+1], a, valid[:, :j+1], torch.ones(a.shape[:-1], dtype=torch.bool, device=a.device))
completed = enc['effect_post'][0][enc['effect_complete'][0]].float().cpu()
history = torch.zeros(4, service._cte.cfg.effect_dim)
mask = torch.zeros(4, dtype=torch.bool)
take = min(4, len(completed))
if take: history[-take:] = completed[-take:]; mask[-take:] = True
for k,v in zip(features, (global_value,enc['phase'][0,-1].float().cpu(),history,mask)):
features[k].append(v)
payload = {k: torch.stack(v) for k,v in features.items()}
payload['episode_index'] = ep
payload['instruction'] = instruction
payload['fingerprint'] = cache_fingerprint
payload['starts'] = torch.tensor(starts, dtype=torch.int64)
temp = path.with_suffix(f'.{os.getpid()}.tmp'); torch.save(payload,temp); temp.replace(path)
print(f'prepared {count+1}/{len(episodes)} episode={ep} windows={len(starts)}',flush=True)
if __name__ == '__main__':
try:
main()
finally:
if torch.distributed.is_initialized():
torch.distributed.destroy_process_group()
|