RTC adjustments. Bug fix & Alex Soare optimization

Open helper2424 opened this issue 1 month ago • 0 comments

What this does

Implement suggestion from the https://alexander-soare.github.io/robotics/2025/08/05/smooth-as-butter-robot-policies.html.

Two important changes:

Added the logic that allows skipping max_guidance param and use number of flow matching steps as basic clipping parameter.
The guidance calucaltion is extended with sigma_d. The default value is 1.0, so the default behavior is equal to roginal paper. But library users can adjust this value
Fix the bug with in-painting algorithm

How it was tested

Run pytest tests/policies/pi0_pi05
Run pytest tests/policies/smolvla
Run pytest tests/policies/rtc
Also, was run the following script to check different params for RTC:

#!/bin/bash

# Script to run RTC evaluation experiments with different parameters
# This script tests various combinations of:
# - num_inference_steps (flow matching steps)
# - sigma_d (variance clipping parameter)
# - Different policies (SmolVLA and PI0.5)

set -e  # Exit on error

# Color codes for output
GREEN='\033[0;32m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration arrays
NUM_STEPS=(2 5 10 20 50 100)
SIGMA_D_VALUES=(0.1 0.2 0.5 0.8 1.0 1.2 1.5)

# Model configurations
SMOLVLA_POLICY="helper2424/smolvla_check_rtc_last3"
SMOLVLA_DATASET="helper2424/check_rtc"
PI05_POLICY="lerobot/pi05_libero_finetuned"
PI05_DATASET="HuggingFaceVLA/libero"

# Common parameters
DEVICE="mps"
SEED=10
EXECUTION_HORIZON=8
ATTENTION_SCHEDULE="EXP"
INFERENCE_DELAY=4

# Create results directory
RESULTS_DIR="rtc_experiments_results"
mkdir -p "$RESULTS_DIR"

# Log file
LOG_FILE="$RESULTS_DIR/experiment_log.txt"
echo "RTC Evaluation Experiments - $(date)" > "$LOG_FILE"
echo "======================================" >> "$LOG_FILE"

# Function to run a single experiment
run_experiment() {
    local model_name=$1
    local policy_path=$2
    local dataset_repo=$3
    local num_steps=$4
    local sigma_d=$5

    local output_dir="${RESULTS_DIR}/${model_name}_steps_${num_steps}_sigma_${sigma_d}"

    echo -e "${BLUE}Running: ${model_name} | steps=${num_steps} | sigma_d=${sigma_d}${NC}"
    echo "$(date): Starting ${model_name} steps=${num_steps} sigma_d=${sigma_d}" >> "$LOG_FILE"

    # Run the evaluation
    uv run python examples/rtc/eval_dataset.py \
        --policy.path="$policy_path" \
        --dataset.repo_id="$dataset_repo" \
        --rtc.execution_horizon="$EXECUTION_HORIZON" \
        --rtc.sigma_d="$sigma_d" \
        --device="$DEVICE" \
        --rtc.prefix_attention_schedule="$ATTENTION_SCHEDULE" \
        --seed="$SEED" \
        --num_inference_steps="$num_steps" \
        --inference_delay="$INFERENCE_DELAY" \
        --output_dir="$output_dir" 2>&1 | tee -a "$LOG_FILE"

    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✓ Completed: ${model_name} | steps=${num_steps} | sigma_d=${sigma_d}${NC}"
        echo "$(date): SUCCESS ${model_name} steps=${num_steps} sigma_d=${sigma_d}" >> "$LOG_FILE"
    else
        echo "ERROR: Failed ${model_name} steps=${num_steps} sigma_d=${sigma_d}" >> "$LOG_FILE"
        echo "Continuing with next experiment..."
    fi

    echo "" >> "$LOG_FILE"
}

# Main execution loop
echo "Starting RTC evaluation experiments..."
echo "Results will be saved to: $RESULTS_DIR"
echo ""

# # Run experiments for SmolVLA
# echo "=========================================="
# echo "Running SmolVLA experiments"
# echo "=========================================="
# for num_steps in "${NUM_STEPS[@]}"; do
#     for sigma_d in "${SIGMA_D_VALUES[@]}"; do
#         run_experiment \
#             "smolvla" \
#             "$SMOLVLA_POLICY" \
#             "$SMOLVLA_DATASET" \
#             "$num_steps" \
#             "$sigma_d"
#     done
# done

# Run experiments for PI0.5
echo "=========================================="
echo "Running PI0.5 experiments"
echo "=========================================="
for num_steps in "${NUM_STEPS[@]}"; do
    for sigma_d in "${SIGMA_D_VALUES[@]}"; do
        run_experiment \
            "pi05" \
            "$PI05_POLICY" \
            "$PI05_DATASET" \
            "$num_steps" \
            "$sigma_d"
    done
done

echo ""
echo "=========================================="
echo "All experiments completed!"
echo "Results saved to: $RESULTS_DIR"
echo "Log file: $LOG_FILE"
echo "=========================================="

# Generate summary
SUMMARY_FILE="$RESULTS_DIR/summary.txt"
echo "Experiment Summary - $(date)" > "$SUMMARY_FILE"
echo "======================================" >> "$SUMMARY_FILE"
echo "" >> "$SUMMARY_FILE"
echo "Total experiments: $(( ${#NUM_STEPS[@]} * ${#SIGMA_D_VALUES[@]} * 2 ))" >> "$SUMMARY_FILE"
echo "Models tested: SmolVLA, PI0.5" >> "$SUMMARY_FILE"
echo "Num steps tested: ${NUM_STEPS[*]}" >> "$SUMMARY_FILE"
echo "Sigma_d values tested: ${SIGMA_D_VALUES[*]}" >> "$SUMMARY_FILE"
echo "" >> "$SUMMARY_FILE"
echo "Results directory structure:" >> "$SUMMARY_FILE"
find "$RESULTS_DIR" -type d -name "*_steps_*" | sort >> "$SUMMARY_FILE"

echo ""
echo "Summary saved to: $SUMMARY_FILE"