| """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) |
| |
| |
| 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() |
|
|
