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