Upload lapa_bridge_3d_all.sh

lapa_bridge_3d_all.sh

@@ -0,0 +1,147 @@
+#!/usr/bin/env bash
+# Full SimplerEnv Bridge suite (4 tasks x 24 episodes) - finetune-matched LAPA config.
+# Intended for product_vq_3d (or any ckpt); same flags as lapa_bridge_smoke.sh.
+#
+# Archives existing ./results and evaluation_results.txt before running.
+#
+#   export CKPT_PATH=params::/path/to/streaming_params_N
+#   bash "${LAPA_ROOT}/SimplerEnv/scripts/lapa_bridge_3d_all.sh"
+#
+# Optional: EPISODE_END=24 (default), ARCHIVE_PREVIOUS=1 (default), SKIP_TASKS=carrot (comma list)
+
+set -euo pipefail
+
+gpu_id=0
+policy_model=lapa
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+SIMPLER_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+LAPA_ROOT="$(cd "${SIMPLER_ROOT}/.." && pwd)"
+
+: "${CKPT_PATH:?Set CKPT_PATH to params::/abs/path/to/streaming_params_*}"
+
+ACTION_SCALE_FILE="${ACTION_SCALE_FILE:-${LAPA_ROOT}/data/bridge_finetune_filtered.csv}"
+EPISODE_START="${EPISODE_START:-0}"
+EPISODE_END="${EPISODE_END:-24}"
+ARCHIVE_PREVIOUS="${ARCHIVE_PREVIOUS:-1}"
+SKIP_TASKS="${SKIP_TASKS:-}"
+
+LLAMA_UPDATE_CFG='dict(action_vocab_size=256,delta_vocab_size=8,sample_mode='"'"'text'"'"',theta=50000000,max_sequence_length=32768,scan_attention=False,scan_query_chunk_size=128,scan_key_chunk_size=128,scan_mlp=False,scan_mlp_chunk_size=8192,scan_layers=True)'
+
+ckpt_dir="${CKPT_PATH#params::}"
+for f in "${ACTION_SCALE_FILE}" "${LAPA_ROOT}/lapa_checkpoints/tokenizer.model" "${LAPA_ROOT}/lapa_checkpoints/vqgan"; do
+  [[ -f "${f}" ]] || { echo "ERROR: missing ${f}"; exit 1; }
+done
+[[ -e "${ckpt_dir}" ]] || { echo "ERROR: checkpoint not found: ${ckpt_dir}"; exit 1; }
+
+_should_skip() {
+  local tag="$1"
+  [[ -z "${SKIP_TASKS}" ]] && return 1
+  case ",${SKIP_TASKS}," in
+    *",${tag},"*) return 0 ;;
+  esac
+  return 1
+}
+
+_archive_previous() {
+  [[ "${ARCHIVE_PREVIOUS}" == "1" ]] || return 0
+  if [[ ! -d "${SIMPLER_ROOT}/results" && ! -f "${SIMPLER_ROOT}/evaluation_results.txt" ]]; then
+    echo "No prior results to archive."
+    return 0
+  fi
+  local stamp
+  stamp="$(date -u +%Y%m%dT%H%M%SZ)"
+  local archive_dir="${SIMPLER_ROOT}/results_archive/${stamp}_pre_full_bridge"
+  mkdir -p "${archive_dir}"
+  if [[ -d "${SIMPLER_ROOT}/results" ]]; then
+    mv "${SIMPLER_ROOT}/results" "${archive_dir}/results"
+    echo "Archived results -> ${archive_dir}/results"
+  fi
+  if [[ -f "${SIMPLER_ROOT}/evaluation_results.txt" ]]; then
+    mv "${SIMPLER_ROOT}/evaluation_results.txt" "${archive_dir}/evaluation_results.txt"
+    echo "Archived evaluation_results.txt -> ${archive_dir}/"
+  fi
+  mkdir -p "${SIMPLER_ROOT}/results"
+}
+
+_run_task() {
+  local tag="$1"
+  local env_name="$2"
+  local scene_name="$3"
+  local robot="$4"
+  local overlay="$5"
+  local max_steps="$6"
+  local robot_x="$7"
+  local robot_y="$8"
+
+  if _should_skip "${tag}"; then
+    echo "=== SKIP task ${tag} (${env_name}) ==="
+    return 0
+  fi
+
+  [[ -f "${overlay}" ]] || { echo "ERROR: missing overlay ${overlay} for ${tag}"; exit 1; }
+
+  echo ""
+  echo "=== LAPA Bridge eval: ${tag} ==="
+  echo "ENV_NAME=${env_name} episodes=[${EPISODE_START},${EPISODE_END}) max_steps=${max_steps}"
+  echo "scene=${scene_name} robot=${robot}"
+  nvidia-smi || true
+
+  CUDA_VISIBLE_DEVICES=${gpu_id} python "${SIMPLER_ROOT}/simpler_env/main_inference_lapa.py" \
+    --policy-model "${policy_model}" \
+    --ckpt-path "${CKPT_PATH}" \
+    --robot "${robot}" \
+    --policy-setup widowx_bridge \
+    --action-scale-file "${ACTION_SCALE_FILE}" \
+    --control-freq 5 \
+    --sim-freq 500 \
+    --max-episode-steps "${max_steps}" \
+    --env-name "${env_name}" \
+    --scene-name "${scene_name}" \
+    --rgb-overlay-path "${overlay}" \
+    --robot-init-x "${robot_x}" "${robot_x}" 1 \
+    --robot-init-y "${robot_y}" "${robot_y}" 1 \
+    --obj-variation-mode episode \
+    --obj-episode-range "${EPISODE_START}" "${EPISODE_END}" \
+    --robot-init-rot-quat-center 0 0 0 1 \
+    --robot-init-rot-rpy-range 0 0 1 0 0 1 0 0 1 \
+    --load-llama-config 7b \
+    --update-llama-config "${LLAMA_UPDATE_CFG}" \
+    --vocab-file "${LAPA_ROOT}/lapa_checkpoints/tokenizer.model" \
+    --vqgan-checkpoint "${LAPA_ROOT}/lapa_checkpoints/vqgan" \
+    --mesh-dim '1,-1,1,1' \
+    --tokens-per-action 7 \
+    --tokens-per-delta 4 \
+    --tf-memory-limit 2048
+}
+
+cd "${SIMPLER_ROOT}"
+_archive_previous
+
+OVERLAY_V1="${SIMPLER_ROOT}/ManiSkill2_real2sim/data/real_inpainting/bridge_real_eval_1.png"
+OVERLAY_SINK="${SIMPLER_ROOT}/ManiSkill2_real2sim/data/real_inpainting/bridge_sink.png"
+
+echo "=== LAPA Bridge full suite (3D / finetune config) ==="
+echo "CKPT_PATH=${CKPT_PATH}"
+echo "ACTION_SCALE_FILE=${ACTION_SCALE_FILE}"
+echo "EPISODE_RANGE=[${EPISODE_START},${EPISODE_END})"
+
+_run_task stack \
+  StackGreenCubeOnYellowCubeBakedTexInScene-v0 \
+  bridge_table_1_v1 widowx "${OVERLAY_V1}" 60 0.147 0.028
+
+_run_task carrot \
+  PutCarrotOnPlateInScene-v0 \
+  bridge_table_1_v1 widowx "${OVERLAY_V1}" 60 0.147 0.028
+
+_run_task spoon \
+  PutSpoonOnTableClothInScene-v0 \
+  bridge_table_1_v1 widowx "${OVERLAY_V1}" 60 0.147 0.028
+
+_run_task eggplant \
+  PutEggplantInBasketScene-v0 \
+  bridge_table_1_v2 widowx_sink_camera_setup "${OVERLAY_SINK}" 120 0.127 0.06
+
+echo ""
+echo "Done. Fresh results under ${SIMPLER_ROOT}/results/"
+echo "Summary appended to ${SIMPLER_ROOT}/evaluation_results.txt"