This folder contains the plotting utilities for the gradient-analysis workflow.

There are two plotting entry points:

1. plot_gradient_analysis.py

• pulls one W&B run directly
• exports local json / csv
• writes per-step PNG plots

2. plot_icml_steps.py

• builds a fixed 3-step comparison figure from already-exported metrics.json files
• intended for paper-style summary figures

For the training-side workflow and arguments, see:

• docs/guide_gradient_analysis.md

Current default training behavior from config/base.yaml:

• trainer.gradient_analysis_mode=True
• trainer.gradient_analysis_every=50
• trainer.gradient_analysis_env_groups=null
• trainer.gradient_analysis_group_size=null
• trainer.exit_after_gradient_analysis=False

Example helper runner:

bash scripts/runs/run_sokoban_ppo_filter_grad_analysis.sh \
  --gpus 0,1,2,3,4,5,6,7

That job:

• trains for 101 steps
• validates before training and every 10 steps
• runs gradient analysis at steps 1, 51, and 101
• uses a training batch of 8x16
• uses a separate gradient-analysis batch of 128x16

python gradient_analysis/plot_gradient_analysis.py \
  --wandb-path deimos-xing/ragen_gradient_analysis/<run_id> \
  --list-steps

python gradient_analysis/plot_gradient_analysis.py \
  --wandb-path deimos-xing/ragen_gradient_analysis/<run_id>

Default output directory:

gradient_analysis_outputs/<run_name>_<run_id>/

python gradient_analysis/plot_gradient_analysis.py \
  --wandb-path deimos-xing/ragen_gradient_analysis/<run_id> \
  --step 1

python gradient_analysis/plot_gradient_analysis.py \
  --wandb-path deimos-xing/ragen_gradient_analysis/<run_id> \
  --step 1 \
  --output-dir gradient_analysis_outputs/my_custom_dir

For each selected step, the script writes:

• gradient_analysis_summary_step_<N>.png
• gradient_analysis_plots_step_<N>.png
• gradient_analysis_loss_plots_step_<N>.png
• gradient_analysis_reward_std_step_<N>.png
• gradient_analysis_normed_grads_step_<N>.png
• gradient_analysis_metrics_step_<N>.json
• gradient_analysis_bucket_rv_table_step_<N>.csv

The metrics.json export is the bridge to the paper-style plotting script.

If you have three exported step directories and want the fixed grid figure:

python gradient_analysis/plot_icml_steps.py \
  --mode ppo \
  --step0-dir /path/to/step0 \
  --step20-dir /path/to/step20 \
  --step40-dir /path/to/step40 \
  --out gradient_analysis_outputs/ppo_step0_20_40.png

Each step directory must contain:

metrics.json

If your exported file is named gradient_analysis_metrics_step_<N>.json, copy or rename it to metrics.json inside each step directory before calling plot_icml_steps.py.

For a new run, start with:

1. gradient_analysis_summary_step_<N>.png
2. gradient_analysis_plots_step_<N>.png
3. gradient_analysis_metrics_step_<N>.json

Those three are usually enough to tell:

• how many buckets were populated
• whether task gradients dominate regularizer gradients
• whether gradient magnitude is monotonic or non-monotonic in reward variance