File size: 2,503 Bytes
952993c | 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 | #!/usr/bin/env bash
# Precompute the T5 text embedding for a task prompt into the PERSISTENT cache
# (fastwam_train/text_embeds_cache/). The dataloader hard-fails on a cache
# miss, so this must exist before training. GPU strongly recommended (loads
# the Wan2.2 T5 encoder); works on CPU but slowly.
#
# bash scripts/precompute_text_embeds.sh ["prompt string"] [task_cfg_name]
# e.g.: bash scripts/precompute_text_embeds.sh \
# "Lift the lid, put it aside, and pick the black plum."
#
# Uses the fastwam precompute script's `override_instruction` path: it encodes
# exactly this one prompt and does NOT scan the dataset's tasks.jsonl, so it can
# run even before the LeRobot conversion exists -- but the string MUST match the
# --task string baked into the dataset's tasks.jsonl, or the training-time hash
# lookup will miss. Writes <sha256(DEFAULT_PROMPT.format(task=...))>.t5_len128.wan22ti2v5b.pt.
set -euo pipefail
TASK_STR="${1:-place the cube in the bowl}"
# Any composable task config works here (override_instruction bypasses its
# actual data/prompt) -- just needs to exist. place_cube_flashwam_ft was
# removed 2026-08-08; pcbnew100_flashwam_scratch is the current stand-in.
TASK_CFG="${2:-pcbnew100_flashwam_scratch}"
REPO=/shared_work/amin-a/realworld-wam
FT="${REPO}/fastwam_train"
CONFIG_DIR="${FT}/configs"
CACHE="${FT}/text_embeds_cache"
FASTWAM_ROOT=/shared_work/physical_intelligence/policies/Fast-WAM/fastwam
PY=/shared_work/environments/miniconda3/envs/fastwam-robotwin-rui/bin/python
mkdir -p "${CACHE}"
cd "${FASTWAM_ROOT}" # read-only: T5/tokenizer weights via relative checkpoints/ paths
# hydra/job_logging=stdout: else Hydra tries to write its log into the read-only
# checkout cwd and crashes at startup (George's env quirk #1).
# data.train.text_embedding_cache_dir override -> write to the PERSISTENT cache
# (the config's default points at /dev/shm, which is a training-time staging dir).
# Hydra's CLI parser treats a bare comma in a value as a list separator (e.g.
# a prompt like "Lift the lid, put it aside, ..." would otherwise be parsed as
# a 2-element list, not a string) -- wrap in escaped single-quotes to force
# string interpretation, per Hydra's own error message.
"${PY}" scripts/precompute_text_embeds.py \
--config-dir "${CONFIG_DIR}" \
"task=${TASK_CFG}" \
"+override_instruction='${TASK_STR}'" \
"data.train.text_embedding_cache_dir=${CACHE}" \
hydra/job_logging=stdout
echo "cache contents:"
ls -la "${CACHE}"
|