p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/mihir/berlin/diffusion/coinrun/sample_coinrun.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/sample/coinrun/diffusion/sample_coinrun_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/sample/coinrun/diffusion/sample_coinrun_%j.log
#SBATCH --job-name=sample_coinrun

# Activate virtual environment
source .venv/bin/activate

array_records_dir=/fast/project/HFMI_ProFOUND/mihir/data/data_coinrun/coinrun_episodes_50m/test

CHECKPOINT_PATH=/fast/project/HFMI_ProFOUND/mihir/checkpoints/dyn/train_dyn_dit/32748
# CHECKPOINT_PATH=/fast/project/HFMI_ProFOUND/mihir/checkpoints/dyn/train_dyn_dit/32772

echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint "$CHECKPOINT_PATH" \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=16 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=diffusion \
    --diffusion_denoise_steps=128 \
    --diffusion_corrupt_context_factor=0.1 \
    --print-action-indices \
    --output_dir "gifs/diffusion/coinrun/lr_1e-4" \
    --tokenizer_type=mae \
[tab] slurm/utils/mihir/plotting/run-renamer.py (python)
import pandas as pd
import sys
import argparse

def rename_columns_in_csv(input_csv, output_csv, column_mapping):
    try:
        # Read the CSV file
        df = pd.read_csv(input_csv)
        
        # Rename the columns using the provided mapping for substrings
        new_columns = {}
        for col in df.columns:
            new_col = col
            for old_substring, new_substring in column_mapping.items():
                if old_substring in col:
                    new_col = new_col.replace(old_substring, new_substring)
            new_columns[col] = new_col
        
        df.rename(columns=new_columns, inplace=True)
        
        # Save the modified DataFrame to a new CSV file
        df.to_csv(output_csv, index=False)
        print(f"Columns renamed and saved to {output_csv}")
    except Exception as e:
        print(f"An error occurred: {e}", file=sys.stderr)

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Rename columns in a CSV file.")
    parser.add_argument("input_csv", help="Path to the input CSV file")
    parser.add_argument("output_csv", help="Path to the output CSV file")
    args = parser.parse_args()

    # Define the column mapping here
    column_mapping = {
        "dynamics_coinrun_mila_submission_no_cotraining_29978": "w/o co-training",
        "dynamics_coinrun_mila_submission_gt_actions_29977": "ground-truth actions",
        "dynamics_coinrun_mila_submission_ffn_dim_ablation_29976": "1x feedforward expansion",
        "dynamics_coinrun_mila_submission_29973": "jasmine-base",
        "dynamics_coinrun_mila_submission_causal_29974": "causal",
        "dynamics_case_study_dataset_10M_30031": "Genie config",
        "dynamics_case_study_dataset_10M_action_prepend_30032": "Genie w/ action-prepend",
        "val_full_frame_psnr": "Rollout PSNR",
        "val_full_frame_ssim": "Rollout SSIM",
        "val_cross_entropy_loss": "Validation loss",
        "cross_entropy_loss": "Train loss"
        # Add more mappings as needed
    }

    rename_columns_in_csv(args.input_csv, args.output_csv, column_mapping)
[tab] slurm/utils/mihir/plotting/causal_ablation/make_plots_multiplot.sh (shellscript)
# python slurm/utils/mihir/run-renamer.py slurm/utils/mihir/causal_ablation/full_frame_psnr_causal.csv slurm/utils/mihir/causal_ablation/full_frame_psnr_causal_renamed.csv
# python slurm/utils/mihir/run-renamer.py slurm/utils/mihir/causal_ablation/ull_frame_psnr_ssim.csv slurm/utils/mihir/causal_ablation/ull_frame_psnr_ssim_renamed.csv
# python slurm/utils/mihir/run-renamer.py slurm/utils/mihir/causal_ablation/causal_val_loss.csv slurm/utils/mihir/causal_ablation/causal_val_loss_renamed.csv

python slurm/utils/mihir/causal_plot.py --metric "Validation loss" "Rollout PSNR" "Rollout SSIM" --x=step --save=slurm/utils/mihir/causal_ablation/multiplot.pdf slurm/utils/mihir/causal_ablation/full_frame_psnr_causal_renamed.csv slurm/utils/mihir/causal_ablation/ull_frame_psnr_ssim_renamed.csv slurm/utils/mihir/causal_ablation/causal_val_loss_renamed.csv
[tab] slurm/utils/mihir/plotting/dit-speedrun/make_plots_multiplot.sh (shellscript)
# python slurm/utils/mihir/run-renamer.py slurm/utils/mihir/causal_ablation/full_frame_psnr_causal.csv slurm/utils/mihir/causal_ablation/full_frame_psnr_causal_renamed.csv
# python slurm/utils/mihir/run-renamer.py slurm/utils/mihir/causal_ablation/ull_frame_psnr_ssim.csv slurm/utils/mihir/causal_ablation/ull_frame_psnr_ssim_renamed.csv
# python slurm/utils/mihir/run-renamer.py slurm/utils/mihir/causal_ablation/causal_val_loss.csv slurm/utils/mihir/causal_ablation/causal_val_loss_renamed.csv

python slurm/utils/mihir/causal_plot.py --metric "Validation loss" "Rollout PSNR" "Rollout SSIM" --x=step --save=slurm/utils/mihir/causal_ablation/multiplot.pdf slurm/utils/mihir/causal_ablation/full_frame_psnr_causal_renamed.csv slurm/utils/mihir/causal_ablation/ull_frame_psnr_ssim_renamed.csv slurm/utils/mihir/causal_ablation/causal_val_loss_renamed.csv
[tab] slurm/utils/mihir/plotting/dit-speedrun/dit-maskgit-compare-psnr-renamed.csv (plaintext)
step,maskgit - _step,maskgit - _step__MIN,maskgit - _step__MAX,maskgit - Rollout PSNR,maskgit - Rollout PSNR__MIN,maskgit - Rollout PSNR__MAX,diffusion - _step,diffusion - _step__MIN,diffusion - _step__MAX,diffusion - Rollout PSNR,diffusion - Rollout PSNR__MIN,diffusion - Rollout PSNR__MAX
20000,418,418,418,24.81561623367609,24.81561623367609,24.81561623367609,418,418,418,26.00116662418141,26.00116662418141,26.00116662418141
40000,838,838,838,26.404946869494868,26.404946869494868,26.404946869494868,838,838,838,28.054943196913776,28.054943196913776,28.054943196913776
60000,1258,1258,1258,26.93376903907925,26.93376903907925,26.93376903907925,1258,1258,1258,28.116503883810605,28.116503883810605,28.116503883810605
80000,1678,1678,1678,27.40945202696557,27.40945202696557,27.40945202696557,1678,1678,1678,28.894807254566867,28.894807254566867,28.894807254566867
100000,2098,2098,2098,27.19209024017932,27.19209024017932,27.19209024017932,2098,2098,2098,28.20263069751216,28.20263069751216,28.20263069751216
120000,2518,2518,2518,27.46689265382056,27.46689265382056,27.46689265382056,2518,2518,2518,28.688626457663144,28.688626457663144,28.688626457663144
140000,2938,2938,2938,27.808968974094764,27.808968974094764,27.808968974094764,2938,2938,2938,29.43852944467582,29.43852944467582,29.43852944467582
160000,3358,3358,3358,27.770290748745783,27.770290748745783,27.770290748745783,3358,3358,3358,29.361815807866115,29.361815807866115,29.361815807866115
180000,3778,3778,3778,27.855009004181507,27.855009004181507,27.855009004181507,3778,3778,3778,29.528736750284832,29.528736750284832,29.528736750284832
200000,4198,4198,4198,27.8867036782059,27.8867036782059,27.8867036782059,4198,4198,4198,29.555948743633195,29.555948743633195,29.555948743633195
[content] slurm/utils/mihir/plotting/wandb_csv_to_matplotlib_multiplot.py (python)
                # Assign colors based on run name
                color = None
                run_name_lower = run_name.lower()
                if "maskgit" in run_name_lower:
                    color = 'C1'  # Classic matplotlib orange
                elif "diffusion" in run_name_lower:
                    color = 'C0'  # Classic matplotlib blue
                
                line, = ax.plot(x, y, label=label, linewidth=3.5, marker="o", markersize=7, color=color)
[tab] slurm/jobs/mihir/berlin/diffusion/doom/sample_coinrun.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/sample/coinrun/diffusion/sample_coinrun_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/sample/coinrun/diffusion/sample_coinrun_%j.log
#SBATCH --job-name=sample_coinrun

# Activate virtual environment
source .venv/bin/activate

array_records_dir=/fast/project/HFMI_ProFOUND/mihir/data/data_coinrun/coinrun_episodes_50m/test

CHECKPOINT_PATH=/fast/project/HFMI_ProFOUND/mihir/checkpoints/dyn/train_dyn_dit_coinrun/33016

echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/baselines/diffusion/sample_diffusion.py \
    --checkpoint "$CHECKPOINT_PATH" \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=16 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --diffusion_denoise_steps=25 \
    --diffusion_corrupt_context_factor=0.1 \
    --print-action-indices \
    --output_dir "gifs/diffusion/coinrun/diffusion/submission" \
[tab] slurm/jobs/mihir/berlin/maskgit/sample_coinrun.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/sample/coinrun/diffusion/sample_coinrun_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/sample/coinrun/diffusion/sample_coinrun_%j.log
#SBATCH --job-name=sample_coinrun

# Activate virtual environment
source .venv/bin/activate

array_records_dir=/fast/project/HFMI_ProFOUND/mihir/data/data_coinrun/coinrun_episodes_50m/test

CHECKPOINT_PATH=/fast/project/HFMI_ProFOUND/mihir/checkpoints/dyn/train_dyn_maskgit_coinrun/33019

echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/baselines/maskgit/sample_maskgit.py \
    --checkpoint "$CHECKPOINT_PATH" \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=16 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --print-action-indices \
    --output_dir "gifs/maskgit/coinrun/submission" \
[content] jasmine/baselines/maskgit/sample_maskgit.py (python)
    # Save predicted videos as PNG image with all frames (skipping first 4) for each item in batch
    skip = 4
    # pred_videos: (B, T, H, W, C)
    B, T, H, W, C = pred_videos.shape

    for i in range(B):
        # Predicted
        pred_strip = np.concatenate(
            [pred_videos[i, t] for t in range(skip, args.seq_len)], axis=1
        )  # resulting shape: (H, (T-skip)*W, C)
        pred_img = Image.fromarray(pred_strip)
        pred_img.save(
            os.path.join(args.output_dir, f"generation-{args.seed}-sample{i}-diffusion.png")
        )
        # GT
        gt_strip = np.concatenate(
            [true_videos[i, t] for t in range(skip, args.seq_len)], axis=1
        )  # resulting shape: (H, (T-skip)*W, C)
        gt_img = Image.fromarray(gt_strip)
        gt_img.save(
            os.path.join(args.output_dir, f"generation-{args.seed}-sample{i}-gt.png")
        )
[content] slurm/utils/mihir/plotting/dit-speedrun/create_sampling_fig.sh (shellscript)

gifs/diffusion/coinrun/diffusion/submission/generation-3-sample1-diffusion.png
gifs/diffusion/coinrun/diffusion/submission/generation-3-sample2-diffusion.png
gifs/diffusion/coinrun/diffusion/submission/generation-3-sample3-diffusion.png
gifs/maskgit/coinrun/submission/generation-3-sample0-maskgit.png
gifs/maskgit/coinrun/submission/generation-3-sample1-maskgit.png
gifs/maskgit/coinrun/submission/generation-3-sample2-maskgit.png
gifs/maskgit/coinrun/submission/generation-3-sample3-maskgit.png
[content] slurm/utils/mihir/plotting/dit-speedrun/create_sampling_fig.sh (shellscript)

    gifs/maskgit/coinrun/submission/generation-3-sample1-maskgit.png
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample2-gt.png
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample2-diffusion.png
    gifs/maskgit/coinrun/submission/generation-3-sample2-maskgit.png
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample3-gt.png
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample3-diffusion.png
    gifs/maskgit/coinrun/submission/generation-3-sample3-maskgit.png
[content] jasmine/baselines/maskgit/sample_maskgit.py (python)

        if args.print_action_indices:
            d = ImageDraw.Draw(pred_img)
            for t in range(skip, args.seq_len):
                if action_batch_E is not None:
                    action = int(action_batch_BSm11[i, t - 1, 0])
                else:
                    action = int(batch["actions"][i, t - 1])
                x_offset = (t - skip) * pred_videos.shape[3] + 2
[tab] slurm/utils/mihir/plotting/dit-speedrun/create_sampling_fig.py (python)
#!/usr/bin/env bash

# List of files to concatenate (edit this list as needed)
files=(
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample0-gt.png
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample0-diffusion.png
    gifs/maskgit/coinrun/submission/generation-3-sample0-maskgit.png
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample1-gt.png
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample1-diffusion.png
    gifs/maskgit/coinrun/submission/generation-3-sample1-maskgit.png
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample2-gt.png
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample2-diffusion.png
    gifs/maskgit/coinrun/submission/generation-3-sample2-maskgit.png
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample3-gt.png
    gifs/diffusion/coinrun/diffusion/submission/generation-3-sample3-diffusion.png
    gifs/maskgit/coinrun/submission/generation-3-sample3-maskgit.png
)

# Output file name
output_file="slurm/utils/mihir/plotting/dit-speedrun/concatenated_results.png"

# Check that all input files exist
for f in "${files[@]}"; do
    if [ ! -f "$f" ]; then
        echo "File not found: $f"
        exit 1
    fi
done

# Concatenate images vertically (under one another) using ImageMagick 'convert'
convert "${files[@]}" -append "$output_file"

echo "Saved concatenated image as: $output_file"