File size: 2,219 Bytes
2622c40 | 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 | #!/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/"
|