diff --git a/scripts/run_dreamwaq_sim2sim.sh b/scripts/run_dreamwaq_sim2sim.sh new file mode 100755 index 0000000..c63aa85 --- /dev/null +++ b/scripts/run_dreamwaq_sim2sim.sh @@ -0,0 +1,124 @@ +#!/usr/bin/env bash +set -Eeuo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(cd -- "${SCRIPT_DIR}/.." && pwd)" +RUNS_DIR="${PROJECT_DIR}/runs/go1-dreamwaq-walk/rslrl" + +checkpoint="" +output="" +terrain="flat" +level="0" +export_only=false + +usage() { + cat <<'EOF' +Export a DreamWaQ checkpoint to ONNX and run MuJoCo sim2sim. + +Usage: + scripts/run_dreamwaq_sim2sim.sh [options] + +Options: + -c, --checkpoint PATH Checkpoint to export. Defaults to the newest model_*.pt. + -o, --output PATH ONNX output path. Defaults to policy.onnx beside checkpoint. + -t, --terrain NAME flat, rough, stairs, dreamwaq, stairs_test, + stairs_box, or flat_stairs. Default: flat. + -l, --level N Terrain difficulty level passed to sim2sim. Default: 0. + --export-only Export and validate ONNX without starting MuJoCo. + -h, --help Show this help. + +Examples: + scripts/run_dreamwaq_sim2sim.sh + scripts/run_dreamwaq_sim2sim.sh -c runs/.../model_1200.pt + scripts/run_dreamwaq_sim2sim.sh -t dreamwaq -l 3 +EOF +} + +die() { + echo "[ERROR] $*" >&2 + exit 1 +} + +while (($#)); do + case "$1" in + -c|--checkpoint) + (($# >= 2)) || die "$1 requires a path" + checkpoint="$2" + shift 2 + ;; + -o|--output) + (($# >= 2)) || die "$1 requires a path" + output="$2" + shift 2 + ;; + -t|--terrain) + (($# >= 2)) || die "$1 requires a terrain name" + terrain="$2" + shift 2 + ;; + -l|--level) + (($# >= 2)) || die "$1 requires an integer" + level="$2" + shift 2 + ;; + --export-only) + export_only=true + shift + ;; + -h|--help) + usage + exit 0 + ;; + *) + die "unknown argument: $1 (use --help)" + ;; + esac +done + +case "$terrain" in + flat|rough|stairs|dreamwaq|stairs_test|stairs_box|flat_stairs) ;; + *) die "unsupported terrain: ${terrain}" ;; +esac +[[ "$level" =~ ^[0-9]+$ ]] || die "level must be a non-negative integer: ${level}" +command -v uv >/dev/null 2>&1 || die "uv is not available in PATH" + +if [[ -z "$checkpoint" ]]; then + [[ -d "$RUNS_DIR" ]] || die "run directory not found: ${RUNS_DIR}" + checkpoint="$({ + find "$RUNS_DIR" -type f -name 'model_*.pt' -printf '%T@ %p\n' + } | sort -nr | sed -n '1{s/^[^ ]* //;p;}')" + [[ -n "$checkpoint" ]] || die "no model_*.pt checkpoint found under ${RUNS_DIR}" +fi + +if [[ "$checkpoint" != /* ]]; then + checkpoint="${PROJECT_DIR}/${checkpoint}" +fi +[[ -f "$checkpoint" ]] || die "checkpoint not found: ${checkpoint}" + +if [[ -z "$output" ]]; then + output="$(dirname -- "$checkpoint")/policy.onnx" +elif [[ "$output" != /* ]]; then + output="${PROJECT_DIR}/${output}" +fi +mkdir -p -- "$(dirname -- "$output")" + +echo "[Pipeline] checkpoint: ${checkpoint}" +echo "[Pipeline] ONNX: ${output}" + +cd -- "$PROJECT_DIR" +uv run scripts/export_dreamwaq_onnx_new.py \ + --checkpoint "$checkpoint" \ + --output "$output" + +[[ -s "$output" ]] || die "ONNX export did not produce a non-empty file: ${output}" +echo "[Pipeline] ONNX export complete" + +if [[ "$export_only" == true ]]; then + exit 0 +fi + +echo "[Pipeline] starting MuJoCo: terrain=${terrain}, level=${level}" +exec uv run scripts/dreamwaq_sim2sim_mujoco.py \ + --onnx "$output" \ + --terrain "$terrain" \ + --level "$level"