#!/usr/bin/env bash
# Precompute the T5 text embedding for the pusht task prompt into
# pusht_train/text_embeds_cache/. Required before training — the dataloader
# hard-fails on a cache miss, and both sbatch scripts assert the file exists.
#
# Writes: 71cb088954da46da4e3cb6c6f73ac812690e9bccf886c7858377c61918b4e3e8
#           .t5_len128.wan22ti2v5b.pt
#   = sha256("A video recorded from a robot's point of view executing the
#             following instruction: push the T block to the target outline")
#
# Two deliberate choices:
#  * override_instruction=... makes the script encode exactly ONE prompt and
#    skip the dataset scan, so this can run before/independently of the
#    LeRobot conversion (the scan would want meta/tasks.jsonl).
#  * data.train.text_embedding_cache_dir is overridden to the SHARED package
#    dir. The config's own value is the node-local /tmp path that
#    stage_pusht_local.sh populates at job time; writing there from the login
#    node would put the file on the wrong host.
#
# hydra/job_logging=stdout + hydra.run.dir: Hydra's default file handler would
# try to write precompute_text_embeds.log into the read-only fastwam checkout
# (our cwd) and die before loading anything. Same fix as run_one.sh.
#
# CUDA_VISIBLE_DEVICES="" forces CPU: this is a single 128-token forward pass,
# it does not need a GPU, and ad-hoc GPU use outside a submitted job is not
# allowed. Takes a few minutes, almost all of it loading the encoder.
set -euo pipefail

FASTWAM_ROOT=/shared_work/physical_intelligence/policies/Fast-WAM/fastwam
PKG=/shared_work/george/real_world_wam/pusht_train
PYTHON=/shared_work/environments/miniconda3/envs/fastwam-robotwin-rui/bin/python
TASK="push the T block to the target outline"

cd "${FASTWAM_ROOT}"   # read-only: T5/tokenizer weights via relative checkpoints/ paths

CUDA_VISIBLE_DEVICES="" "${PYTHON}" scripts/precompute_text_embeds.py \
  --config-dir "${PKG}/configs" \
  task=pusht_flashwam_scratch \
  "+override_instruction='${TASK}'" \
  "data.train.text_embedding_cache_dir=${PKG}/text_embeds_cache" \
  hydra/job_logging=stdout \
  "hydra.run.dir=${PKG}/.hydra_precompute" \
  "$@"

echo "cache contents:"
ls -la "${PKG}/text_embeds_cache/"