diff --git a/CLAUDE.md b/CLAUDE.md
new file mode 100644
index 0000000..d4d74ab
--- /dev/null
+++ b/CLAUDE.md
@@ -0,0 +1,22 @@
+# Claude Code Instructions
+
+## Python Environment
+
+**ALWAYS use `.venv-310` virtual environment for this project.**
+
+When running Python commands, use:
+```bash
+source .venv-310/bin/activate && python <command>
+```
+
+Never use system `python` or `python3` directly without activating the virtual environment first.
+
+## Example Commands
+
+```bash
+# Running training
+source .venv-310/bin/activate && python -m collab_env.gnn.interaction_particles.run_training --help
+
+# Running any Python script
+source .venv-310/bin/activate && python <script>.py
+```
diff --git a/Dockerfile.gnn b/Dockerfile.gnn
new file mode 100644
index 0000000..4adaa6d
--- /dev/null
+++ b/Dockerfile.gnn
@@ -0,0 +1,13 @@
+FROM runpod/pytorch:2.8.0-py3.11-cuda12.8.1-cudnn-devel-ubuntu22.04
+
+RUN apt-get update && apt-get install -y  \
+    screen \
+    htop \
+    mc \
+    && rm -rf /var/lib/apt/lists/* && \
+    curl https://rclone.org/install.sh | bash && \
+    rclone config create collab-data "google cloud storage" service_account_file=/workspace/collab-environment/config-local/collab-data-463313-c340ad86b28e.json && \
+    pip install uv nvitop && \
+    git config --global user.email "dima.batenkov@gmail.com" && \
+    git config --global user.name "Dmitry Batenkov" 
+WORKDIR /workspace
\ No newline at end of file
diff --git a/collab_env/gnn/interaction_particles/README.md b/collab_env/gnn/interaction_particles/README.md
new file mode 100644
index 0000000..0ac9c7a
--- /dev/null
+++ b/collab_env/gnn/interaction_particles/README.md
@@ -0,0 +1,470 @@
+# InteractionParticle Model for 2D Boids
+
+## Overview
+
+This implementation learns particle interaction dynamics from 2D boid simulation data using a Graph Neural Network (GNN). The model predicts forces between particles based on their relative positions and velocities.
+
+## Model Architecture
+
+### Input Features (per edge)
+
+The model uses **unnormalized** edge features to capture pairwise interactions:
+
+For an edge from particle `j` to particle `i`:
+- **delta_pos**: Relative position `pos_j - pos_i` (2D vector)
+- **r**: Distance `||delta_pos||` (scalar)
+- **delta_vel**: Relative velocity `vel_j - vel_i` (2D vector)
+- **embedding_i**: Learnable particle embedding (16D vector by default)
+
+**Total input size**: 5 + embedding_dim = 21 (with default embedding_dim=16)
+
+### Output
+
+- **Force**: 2D force vector from particle `j` acting on particle `i`
+
+### Architecture Details
+
+- **Message Passing**: Each edge computes a force using an MLP
+- **Aggregation**: Forces from all neighbors are summed (default: 'add')
+- **MLP**: 3-layer fully connected network with hidden_dim=128
+- **Learnable Embeddings**: Each particle has a unique embedding vector
+
+## Key Design Decisions
+
+### 1. No Normalization
+
+**Rationale**: Boid forces have physical units and magnitudes that matter. Normalization would lose this information. The model learns to handle raw physical quantities directly.
+
+### 2. Relative Velocity (not individual velocities)
+
+**Previous**: Used `vel_i` and `vel_j` separately (4 features)
+**Current**: Use `delta_vel = vel_j - vel_i` (2 features)
+
+**Rationale**: Boid rules depend on relative motion, not absolute velocities. This:
+- Reduces redundancy
+- Makes the model Galilean invariant
+- Matches the physics better
+
+### 3. Force Decomposition
+
+Forces can be decomposed into position-dependent and velocity-dependent components:
+
+1. **Total Force**: `F(delta_pos, delta_vel)`
+2. **Position-Only Force**: `F(delta_pos, 0)` - separation & cohesion
+3. **Velocity Residual**: `F(delta_pos, delta_vel) - F(delta_pos, 0)` - alignment
+
+## Visualization Approach
+
+### Setup
+
+- **Node i**: At origin (0, 0), stationary (`vel_i = 0`)
+- **Node j**: At grid position `(x, y)`
+
+### Two Velocity Slices
+
+1. **Away**: `vel_j = (x, y) / ||(x, y)||` (moving away from i)
+2. **Towards**: `vel_j = -(x, y) / ||(x, y)||` (moving towards i)
+
+### Three Force Components (per slice)
+
+For each velocity slice, we plot:
+
+1. **I) Total Force**: Full model output with velocities
+2. **II) Position-Only Force**: Model output with `vel=0` artificially set
+3. **III) Velocity Residual**: Difference `I - II`
+
+This decomposition reveals:
+- How much force comes from relative position (separation/cohesion)
+- How much comes from relative velocity (alignment)
+- Whether the learned model separates these components like true boids
+
+## Quick Start
+
+Train on existing 2D boids data:
+
+```bash
+# Run the quick start script (10 epochs)
+./collab_env/gnn/interaction_particles/train_2d_boids.sh
+
+# Or call run_training directly
+python -m collab_env.gnn.interaction_particles.run_training \
+    --dataset simulated_data/boid_single_species_basic.pt \
+    --epochs 100
+```
+
+Available 2D datasets:
+- `boid_single_species_basic.pt` - Clean flocking behavior
+- `boid_single_species_noisy.pt` - Noisy trajectories
+- `boid_single_species_high_cluster_high_speed.pt` - High clustering
+- `boid_single_species_short.pt` - Short trajectories
+
+## Model Overview
+
+The InteractionParticle model is based on the work from:
+- Paper: [Interaction Networks for Learning about Objects, Relations and Physics (NeurIPS 2016)](https://proceedings.neurips.cc/paper/2016/hash/3147da8ab4a0437c15ef51a5cc7f2dc4-Abstract.html)
+- Implementation adapted from: https://github.com/saalfeldlab/decomp-gnn
+
+### Key Features
+
+1. **Message Passing Architecture**: Uses PyTorch Geometric's message passing framework to learn particle interactions
+2. **Learnable Interaction Functions**: MLPs learn the force functions based on:
+   - Relative positions between particles
+   - Relative velocities
+   - Distance between particles
+   - Learnable particle embeddings
+3. **Graph Construction**: Builds dynamic graphs based on visual range (neighborhood distance)
+
+### Model Components
+
+- **Edge MLP** (`lin_edge`): Multi-layer perceptron that learns the interaction function
+- **Particle Embeddings** (`a`): Learnable embeddings for each particle type
+- **Message Function**: Computes interactions based on relative positions, velocities, and distances
+- **Aggregation**: Sums messages from all neighbors to compute total force
+
+## Installation
+
+The module is part of the `collab_env` package. Ensure you have the following dependencies:
+
+```bash
+pip install torch torch_geometric loguru pyyaml matplotlib scipy
+```
+
+Or install the full environment:
+```bash
+pip install -e .
+```
+
+## Usage
+
+### Quick Start Script
+
+The easiest way to get started:
+
+```bash
+# Edit train_2d_boids.sh to uncomment the command you want
+./collab_env/gnn/interaction_particles/train_2d_boids.sh
+```
+
+The script includes examples for:
+- Quick test (10 epochs)
+- Full training (100 epochs)
+- Different datasets (noisy, high_cluster, short)
+- High capacity models
+- Custom parameters
+- Evaluation only
+
+### Direct Training
+
+Call run_training.py directly:
+
+```bash
+python -m collab_env.gnn.interaction_particles.run_training \
+    --dataset simulated_data/boid_single_species_basic.pt \
+    --config simulated_data/boid_single_species_basic_config.pt \
+    --epochs 100 \
+    --batch-size 32 \
+    --visual-range 0.104 \
+    --save-dir trained_models/my_model
+```
+
+### Command Line Arguments
+
+**Data Arguments:**
+- `--dataset`: Path to boids dataset (.pt file)
+- `--config`: Path to boids config.yaml (for comparison with true rules)
+
+**Training Arguments:**
+- `--epochs`: Number of training epochs (default: 100)
+- `--batch-size`: Batch size (default: 32)
+- `--learning-rate`: Learning rate (default: 1e-3)
+- `--train-split`: Train/validation split ratio (default: 0.8)
+- `--visual-range`: Visual range for edge construction in normalized units (default: 0.3)
+- `--seed`: Random seed (default: 42)
+
+**Model Arguments:**
+- `--hidden-dim`: Hidden dimension for MLP (default: 128)
+- `--embedding-dim`: Particle embedding dimension (default: 16)
+- `--n-layers`: Number of MLP layers (default: 3)
+- `--n-particles`: Number of particles in the dataset (default: 20)
+
+**Output Arguments:**
+- `--save-dir`: Directory to save models and plots
+- `--no-plot`: Skip generating plots
+- `--device`: Device to use (cpu/cuda/cuda:0)
+
+**Evaluation:**
+- `--eval-only`: Only evaluate a saved model
+- `--model-path`: Path to saved model checkpoint for evaluation
+
+### Examples
+
+See `train_2d_boids.sh` for ready-to-run examples, or run directly:
+
+**Quick Test:**
+```bash
+python -m collab_env.gnn.interaction_particles.run_training \
+    --dataset simulated_data/boid_single_species_basic.pt \
+    --epochs 10 \
+    --visual-range 0.104
+```
+
+**Different Datasets:**
+```bash
+# Noisy trajectories
+python -m collab_env.gnn.interaction_particles.run_training \
+    --dataset simulated_data/boid_single_species_noisy.pt \
+    --epochs 100
+
+# High clustering
+python -m collab_env.gnn.interaction_particles.run_training \
+    --dataset simulated_data/boid_single_species_high_cluster_high_speed.pt \
+    --epochs 100
+```
+
+**High Capacity Model:**
+```bash
+python -m collab_env.gnn.interaction_particles.run_training \
+    --hidden-dim 256 \
+    --embedding-dim 32 \
+    --n-layers 4 \
+    --epochs 200
+```
+
+**Evaluate Saved Model:**
+```bash
+python -m collab_env.gnn.interaction_particles.run_training \
+    --eval-only \
+    --model-path trained_models/interaction_particle_2d_*/best_model.pt
+```
+
+## Output
+
+The training script generates the following outputs in `--save-dir`:
+
+### Model Checkpoints
+- `best_model.pt`: Best model (lowest validation loss)
+- `final_model.pt`: Final model after all epochs
+
+### Plots
+- `training_history.png`: Training and validation loss curves
+- `interaction_functions.png`: Learned interaction functions (force vs distance)
+- `comparison_with_boids.png`: Comprehensive comparison with true boid rules
+
+### Logs and Metadata
+- `training.log`: Detailed training log
+- `model_info.yaml`: Model configuration and final metrics
+
+## Comparison with True Boid Rules
+
+The module compares learned interaction functions with the true 2D boid rules implemented in `collab_env/sim/boids_gnn_temp/boid.py`.
+
+### 2D Boid Rules
+
+1. **Separation** (`avoid_others`, line 144-160):
+   - Distance threshold: `min_distance` = 15 pixels
+   - Force: `avoid_factor * (self_pos - other_pos)` = 0.05 * displacement
+   - Effect: **Linear repulsion** at close range
+
+2. **Alignment** (`match_velocity`, line 163-184):
+   - Distance threshold: `visual_range` = 50 pixels
+   - Force: `matching_factor * (avg_velocity - self_velocity)` = 0.5 * vel_diff
+   - Effect: **Step function** - match velocity with neighbors in visual range
+
+3. **Cohesion** (`fly_towards_center`, line 120-141):
+   - Distance threshold: `visual_range` = 50 pixels
+   - Force: `centering_factor * (center_of_mass - self_pos)` = 0.005 * displacement
+   - Effect: **Linear attraction** towards center of neighbors
+
+### Comparison Visualization
+
+The `comparison_with_boids.png` plot shows:
+- **Top Row**: Learned interaction force magnitude and components
+- **Middle Row**: True separation and alignment rules
+- **Bottom Row**: True cohesion rule and normalized comparison
+
+The model should learn:
+- **Short range**: Strong repulsive forces (separation)
+- **Medium range**: Attractive forces (cohesion) and alignment
+- **Long range**: Diminishing interactions
+
+## Data Format
+
+The model expects 2D boids data in `AnimalTrajectoryDataset` format:
+
+```python
+# Each sample is a tuple: (positions, species)
+positions: torch.Tensor  # Shape: [steps, N, 2]
+species: torch.Tensor    # Shape: [N]
+
+# Positions are normalized to [0, 1] by scene size (480x480)
+# steps = number of frames (e.g., 800)
+# N = number of particles (e.g., 20)
+```
+
+The training script automatically:
+- Loads all samples from the dataset
+- Computes velocities and accelerations using finite differences
+- Builds graphs based on visual range
+
+## Model Architecture Details
+
+### Node Features
+For each particle:
+```
+[particle_id, pos_x, pos_y, vel_x, vel_y]
+```
+
+### Edge Features (Message Function)
+For each edge from particle j to i:
+```python
+[
+    delta_pos_x, delta_pos_y,  # Normalized relative position
+    distance,                   # Normalized distance
+    vel_i_x, vel_i_y,          # Normalized velocity of i
+    vel_j_x, vel_j_y,          # Normalized velocity of j
+    embedding_i                 # Learnable particle embedding
+]
+```
+
+### Output
+Predicted 2D acceleration for each particle:
+```
+[acc_x, acc_y]
+```
+
+## Rollout Evaluation
+
+The model supports multi-step autoregressive rollout evaluation to test its ability to predict trajectories over time:
+
+```bash
+# Train with rollout evaluation
+python -m collab_env.gnn.interaction_particles.run_training \
+    --dataset simulated_data/boid_single_species_basic.pt \
+    --epochs 100 \
+    --evaluate-rollout \
+    --n-rollout-steps 50 \
+    --save-dir trained_models/with_rollout
+```
+
+### What is Rollout?
+
+Instead of predicting just one step ahead, rollout uses the model's own predictions as input for the next step, generating entire trajectories. This tests whether the model has learned dynamics that are stable and accurate over time.
+
+### Rollout Outputs
+
+When `--evaluate-rollout` is enabled, the following are generated in `rollout_evaluation/`:
+
+1. **Side-by-side trajectory comparisons** (`rollout_comparison_*.png`):
+   - Left: Ground truth trajectories
+   - Right: Model predicted trajectories
+   - Shows first 3 validation samples
+   - Green markers = start, Red markers = end
+
+2. **Error over time** (`rollout_error_over_time.png`):
+   - Position error as a function of time step
+   - Shows mean and ±1 standard deviation across all validation trajectories
+   - Helps identify if error accumulates over time
+
+3. **Metrics** (logged to console):
+   - Mean position error: Average Euclidean distance between predicted and true positions
+   - Mean velocity error: Average velocity deviation
+   - Statistics computed over all validation trajectories
+
+### Expected Results
+
+For well-trained models on 2D boids:
+- Position error should remain relatively stable over ~50 timesteps
+- Trajectories should maintain flock structure (not drift apart or collapse)
+- Mean position error typically < 0.05 for first 20 steps, then gradually increases
+
+### Validation Dataset
+
+Rollout evaluation uses the **validation split** (last 20% of the dataset by default). This ensures the model is tested on trajectories it hasn't seen during training.
+
+The split matches the training validation split (80/20 by default, configurable via `--train-split`).
+
+## API Usage
+
+You can also use the module programmatically:
+
+```python
+from collab_env.gnn.interaction_particles import (
+    InteractionParticle,
+    train_interaction_particle,
+    evaluate_model,
+    evaluate_rollout,
+    generate_rollout,
+    plot_interaction_functions,
+    compare_with_true_boids,
+    create_rollout_report
+)
+
+# Training
+model, history = train_interaction_particle(
+    dataset_path='collab_env/data/boids/boid_single_species_basic.pt',
+    epochs=100,
+    batch_size=32,
+    save_dir='trained_models/my_model'
+)
+
+# Plotting
+plot_interaction_functions(model, save_path='interaction.png')
+compare_with_true_boids(model, save_path='comparison.png')
+
+# One-step evaluation
+metrics = evaluate_model(model, 'test_data.pt')
+print(f"Test RMSE: {metrics['rmse']}")
+
+# Multi-step rollout evaluation
+rollout_results = evaluate_rollout(
+    model,
+    dataset_path='test_data.pt',
+    visual_range=0.104,
+    n_rollout_steps=50
+)
+print(f"Mean position error: {rollout_results['metrics']['mean_position_error']:.6f}")
+
+# Create rollout visualizations
+create_rollout_report(rollout_results, save_dir='rollout_results')
+
+# Generate single rollout
+import torch
+initial_pos = torch.randn(20, 2)  # 20 particles
+initial_vel = torch.randn(20, 2)
+pred_pos, pred_vel = generate_rollout(model, initial_pos, initial_vel, n_steps=100)
+```
+
+## Performance Tips
+
+1. **Visual Range**: Adjust `--visual-range` based on your data normalization. Too small = sparse graphs, too large = dense graphs
+2. **Batch Size**: Larger batches can speed up training but require more memory
+3. **Model Capacity**: Increase `--hidden-dim` and `--n-layers` for more complex interaction functions
+4. **Training Time**: Expect ~2-5 minutes per epoch on GPU for typical boids datasets
+
+## Troubleshooting
+
+**Issue**: Model not learning / high loss
+- Check data normalization
+- Try increasing model capacity (`--hidden-dim 256`)
+- Adjust visual range (`--visual-range 0.5`)
+- Increase training epochs
+
+**Issue**: Out of memory
+- Reduce `--batch-size`
+- Reduce `--visual-range` (fewer edges)
+- Use CPU (`--device cpu`)
+
+**Issue**: Learned functions don't match boid rules
+- Check that boids config matches the data generation parameters
+- Ensure sufficient training epochs
+- Verify data quality (no NaNs, proper normalization)
+
+## References
+
+1. Battaglia, P. W., et al. (2016). Interaction networks for learning about objects, relations and physics. NeurIPS.
+2. Reynolds, C. W. (1987). Flocks, herds and schools: A distributed behavioral model. SIGGRAPH.
+3. PyTorch Geometric Documentation: https://pytorch-geometric.readthedocs.io/
+
+## License
+
+This code is part of the `collab_env` package and follows the same license (Apache-2.0).
diff --git a/collab_env/gnn/interaction_particles/__init__.py b/collab_env/gnn/interaction_particles/__init__.py
new file mode 100644
index 0000000..1b5872f
--- /dev/null
+++ b/collab_env/gnn/interaction_particles/__init__.py
@@ -0,0 +1,30 @@
+"""
+InteractionParticle model training for 2D boids data.
+Adapted from https://github.com/saalfeldlab/decomp-gnn
+
+Quick start:
+    ./collab_env/gnn/interaction_particles/train_2d_boids.sh
+"""
+
+from .model import InteractionParticle
+from .train import train_interaction_particle, evaluate_model, evaluate_rollout, generate_rollout
+from .plotting import (
+    plot_interaction_functions,
+    compare_with_true_boids,
+    plot_rollout_comparison,
+    plot_rollout_error_over_time,
+    create_rollout_report
+)
+
+__all__ = [
+    'InteractionParticle',
+    'train_interaction_particle',
+    'evaluate_model',
+    'evaluate_rollout',
+    'generate_rollout',
+    'plot_interaction_functions',
+    'compare_with_true_boids',
+    'plot_rollout_comparison',
+    'plot_rollout_error_over_time',
+    'create_rollout_report'
+]
diff --git a/collab_env/gnn/interaction_particles/analyze_all.sh b/collab_env/gnn/interaction_particles/analyze_all.sh
new file mode 100755
index 0000000..b69068b
--- /dev/null
+++ b/collab_env/gnn/interaction_particles/analyze_all.sh
@@ -0,0 +1,7 @@
+#!/bin/bash
+
+python collab_env/gnn/interaction_particles/analyze_dataset.py simulated_data/runpod/boid_single_species_basic.pt --save-dir analysis/runpod_basic
+
+python collab_env/gnn/interaction_particles/analyze_dataset.py simulated_data/runpod/boid_single_species_weakalignment_large.pt --save-dir analysis/runpod_weakalignment
+
+python collab_env/gnn/interaction_particles/analyze_dataset.py simulated_data/runpod/boid_single_species_independent.pt --save-dir analysis/runpod_independent
\ No newline at end of file
diff --git a/collab_env/gnn/interaction_particles/analyze_dataset.py b/collab_env/gnn/interaction_particles/analyze_dataset.py
new file mode 100755
index 0000000..6aa9267
--- /dev/null
+++ b/collab_env/gnn/interaction_particles/analyze_dataset.py
@@ -0,0 +1,516 @@
+#!/usr/bin/env python
+"""
+Exploratory Data Analysis (EDA) for InteractionParticle datasets.
+
+This script analyzes the distribution of input features and output accelerations
+in boid trajectory datasets to understand the data characteristics.
+
+Generates:
+    - eda_distributions.png: 3x3 grid showing distributions of positions, velocities,
+      distances, relative positions/velocities, and accelerations
+    - true_boid_force_fields.png: 2x3 grid showing ground truth boid force decomposition
+      (if config file is available)
+
+Note:
+    For symmetric force decomposition (position-dependent vs velocity-dependent),
+    see run_training.py which generates learned_symmetric_decomposition.png after
+    training a model. This clean separation is most useful for learned models.
+
+Usage:
+    python analyze_dataset.py <dataset_path> [--save-dir <output_dir>]
+
+Example:
+    python analyze_dataset.py simulated_data/boid_single_species_basic.pt --save-dir analysis_output
+"""
+
+import argparse
+import os
+import sys
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+import seaborn as sns
+from pathlib import Path
+from loguru import logger
+
+# Add parent directory to path
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+
+def compute_pairwise_features(positions, velocities):
+    """
+    Compute pairwise features (relative positions, distances, relative velocities).
+
+    Parameters
+    ----------
+    positions : np.ndarray
+        Positions [N, 2]
+    velocities : np.ndarray
+        Velocities [N, 2]
+
+    Returns
+    -------
+    dict
+        Dictionary with 'delta_pos', 'distances', 'delta_vel', 'pos_i', 'pos_j'
+    """
+    N = positions.shape[0]
+
+    # Compute all pairwise features
+    delta_pos = positions[None, :, :] - positions[:, None, :]  # [N, N, 2]
+    distances = np.linalg.norm(delta_pos, axis=-1)  # [N, N]
+    delta_vel = velocities[None, :, :] - velocities[:, None, :]  # [N, N, 2]
+
+    # Flatten and remove self-interactions (diagonal)
+    mask = ~np.eye(N, dtype=bool)
+
+    delta_pos_flat = delta_pos[mask]  # [N*(N-1), 2]
+    distances_flat = distances[mask]  # [N*(N-1)]
+    delta_vel_flat = delta_vel[mask]  # [N*(N-1), 2]
+
+    # Also get absolute positions for both particles
+    # Create indices for i and j
+    i_indices = np.repeat(np.arange(N), N)  # [0,0,...,1,1,...,N-1,N-1,...]
+    j_indices = np.tile(np.arange(N), N)    # [0,1,...,N-1,0,1,...,N-1]
+
+    # Remove self-interactions
+    mask_flat = mask.flatten()
+    i_indices_filtered = i_indices[mask_flat]
+    j_indices_filtered = j_indices[mask_flat]
+
+    pos_i = positions[i_indices_filtered]  # [N*(N-1), 2]
+    pos_j = positions[j_indices_filtered]  # [N*(N-1), 2]
+
+    return {
+        'delta_pos': delta_pos_flat,
+        'distances': distances_flat,
+        'delta_vel': delta_vel_flat,
+        'pos_i': pos_i,
+        'pos_j': pos_j
+    }
+
+
+def analyze_dataset(dataset_path, save_dir=None, device='cpu', scene_size=480.0):
+    """
+    Perform exploratory data analysis on a boid trajectory dataset.
+
+    Parameters
+    ----------
+    dataset_path : str
+        Path to dataset .pt file
+    save_dir : str, optional
+        Directory to save analysis plots
+    device : str, optional
+        Device to use for computations (default: 'cpu')
+    scene_size : float, optional
+        Scene size in pixels for normalizing config parameters (default: 480.0)
+
+    Notes
+    -----
+    Uses natural timestep dt=1.0 from simulation for velocity/acceleration computation.
+    """
+    logger.info(f"Loading dataset from {dataset_path}")
+    logger.info(f"Using device: {device}")
+
+    # Load dataset
+    dataset = torch.load(dataset_path, weights_only=False, map_location=device)
+    logger.info(f"Dataset type: {type(dataset)}")
+    logger.info(f"Number of samples: {len(dataset)}")
+
+    # Extract first sample to check structure
+    positions, species = dataset[0]
+    logger.info(f"Sample shape: positions={positions.shape}, species={species.shape}")
+    T, N, D = positions.shape
+    logger.info(f"Timesteps (T): {T}, Particles (N): {N}, Dimensions (D): {D}")
+
+    # Create save directory
+    if save_dir:
+        os.makedirs(save_dir, exist_ok=True)
+        logger.info(f"Saving analysis to {save_dir}")
+
+    # Collect all data
+    all_positions = []
+    all_velocities = []
+    all_accelerations = []
+    all_delta_pos = []
+    all_distances = []
+    all_delta_vel = []
+    all_pos_i = []
+    all_pos_j = []
+
+    logger.info("Computing features from all samples...")
+    for idx, (positions_traj, _) in enumerate(dataset):
+        if idx % 10 == 0:
+            logger.info(f"Processing sample {idx}/{len(dataset)}")
+
+        positions_np = positions_traj.numpy()  # [T, N, 2]
+
+        # Compute velocities (finite differences, natural timestep dt=1.0)
+        # v[t] = p[t+1] - p[t]
+        velocities_np = np.diff(positions_np, axis=0)  # [T-1, N, 2]
+        velocities_np = np.concatenate([velocities_np, velocities_np[-1:]], axis=0)  # [T, N, 2]
+
+        # Compute accelerations (natural timestep dt=1.0)
+        # a[t] = v[t+1] - v[t]
+        accelerations_np = np.diff(velocities_np, axis=0)  # [T-1, N, 2]
+
+        # Collect frame-by-frame
+        for t in range(T-1):  # -1 because acceleration is one shorter
+            pos_t = positions_np[t]  # [N, 2]
+            vel_t = velocities_np[t]  # [N, 2]
+            acc_t = accelerations_np[t]  # [N, 2]
+
+            all_positions.append(pos_t)
+            all_velocities.append(vel_t)
+            all_accelerations.append(acc_t)
+
+            # Compute pairwise features
+            pairwise = compute_pairwise_features(pos_t, vel_t)
+            all_delta_pos.append(pairwise['delta_pos'])
+            all_distances.append(pairwise['distances'])
+            all_delta_vel.append(pairwise['delta_vel'])
+            all_pos_i.append(pairwise['pos_i'])
+            all_pos_j.append(pairwise['pos_j'])
+
+    # Concatenate all data
+    positions_all = np.concatenate(all_positions, axis=0)  # [num_frames*N, 2]
+    velocities_all = np.concatenate(all_velocities, axis=0)
+    accelerations_all = np.concatenate(all_accelerations, axis=0)
+    delta_pos_all = np.concatenate(all_delta_pos, axis=0)  # [num_pairs, 2]
+    distances_all = np.concatenate(all_distances, axis=0)  # [num_pairs]
+    delta_vel_all = np.concatenate(all_delta_vel, axis=0)  # [num_pairs, 2]
+    pos_i_all = np.concatenate(all_pos_i, axis=0)
+    pos_j_all = np.concatenate(all_pos_j, axis=0)
+
+    logger.info(f"Total particles-frames: {len(positions_all)}")
+    logger.info(f"Total pairs: {len(distances_all)}")
+
+    # Print statistics
+    print("\n" + "="*80)
+    print("DATASET STATISTICS")
+    print("="*80)
+
+    print(f"\nPositions (absolute):")
+    print(f"  Range: [{positions_all.min():.4f}, {positions_all.max():.4f}]")
+    print(f"  Mean: {positions_all.mean(axis=0)}")
+    print(f"  Std: {positions_all.std(axis=0)}")
+
+    print(f"\nVelocities (absolute):")
+    print(f"  Range: [{velocities_all.min():.4f}, {velocities_all.max():.4f}]")
+    print(f"  Mean: {velocities_all.mean(axis=0)}")
+    print(f"  Std: {velocities_all.std(axis=0)}")
+    print(f"  Speed mean: {np.linalg.norm(velocities_all, axis=1).mean():.4f}")
+    print(f"  Speed std: {np.linalg.norm(velocities_all, axis=1).std():.4f}")
+
+    print(f"\nAccelerations (targets):")
+    print(f"  Range: [{accelerations_all.min():.4f}, {accelerations_all.max():.4f}]")
+    print(f"  Mean: {accelerations_all.mean(axis=0)}")
+    print(f"  Std: {accelerations_all.std(axis=0)}")
+    print(f"  Magnitude mean: {np.linalg.norm(accelerations_all, axis=1).mean():.4f}")
+    print(f"  Magnitude std: {np.linalg.norm(accelerations_all, axis=1).std():.4f}")
+
+    print(f"\nRelative positions (delta_pos):")
+    print(f"  Range: [{delta_pos_all.min():.4f}, {delta_pos_all.max():.4f}]")
+    print(f"  Mean: {delta_pos_all.mean(axis=0)}")
+    print(f"  Std: {delta_pos_all.std(axis=0)}")
+
+    print(f"\nDistances (between particles):")
+    print(f"  Range: [{distances_all.min():.4f}, {distances_all.max():.4f}]")
+    print(f"  Mean: {distances_all.mean():.4f}")
+    print(f"  Std: {distances_all.std():.4f}")
+    print(f"  Median: {np.median(distances_all):.4f}")
+
+    print(f"\nRelative velocities (delta_vel):")
+    print(f"  Range: [{delta_vel_all.min():.4f}, {delta_vel_all.max():.4f}]")
+    print(f"  Mean: {delta_vel_all.mean(axis=0)}")
+    print(f"  Std: {delta_vel_all.std(axis=0)}")
+    print(f"  Magnitude mean: {np.linalg.norm(delta_vel_all, axis=1).mean():.4f}")
+    print(f"  Magnitude std: {np.linalg.norm(delta_vel_all, axis=1).std():.4f}")
+
+    print("="*80 + "\n")
+
+    # Create visualizations
+    logger.info("Generating visualizations...")
+
+    fig, axes = plt.subplots(3, 3, figsize=(18, 15))
+    fig.suptitle(f'Dataset EDA: {os.path.basename(dataset_path)}', fontsize=16, fontweight='bold')
+
+    # Row 1: Velocities
+    ax = axes[0, 0]
+    ax.scatter(positions_all[:, 0], positions_all[:, 1], alpha=0.1, s=1)
+    ax.set_xlabel('X Position')
+    ax.set_ylabel('Y Position')
+    ax.set_title('Absolute Positions')
+    ax.set_aspect('equal')
+    ax.grid(True, alpha=0.3)
+
+    ax = axes[0, 1]
+    ax.hist2d(velocities_all[:, 0], velocities_all[:, 1], bins=50, cmap='viridis', cmin=1)
+    ax.set_xlabel('Velocity X')
+    ax.set_ylabel('Velocity Y')
+    ax.set_title('Velocity Vector Distribution')
+    ax.set_aspect('equal')
+    plt.colorbar(ax.collections[0], ax=ax, label='Count')
+
+    ax = axes[0, 2]
+    speeds = np.linalg.norm(velocities_all, axis=1)
+    ax.hist(speeds, bins=50, alpha=0.7, color='green')
+    ax.set_xlabel('Speed')
+    ax.set_ylabel('Frequency')
+    ax.set_title('Speed Distribution')
+    ax.grid(True, alpha=0.3)
+    ax.axvline(speeds.mean(), color='red', linestyle='--', label=f'Mean: {speeds.mean():.3f}')
+    ax.legend()
+
+    # Row 2: Pairwise features
+    ax = axes[1, 0]
+    ax.hist(distances_all, bins=50, alpha=0.7, color='purple')
+    ax.set_xlabel('Distance')
+    ax.set_ylabel('Frequency')
+    ax.set_title('Pairwise Distance Distribution')
+    ax.grid(True, alpha=0.3)
+    ax.axvline(distances_all.mean(), color='red', linestyle='--', label=f'Mean: {distances_all.mean():.3f}')
+    ax.legend()
+
+    ax = axes[1, 1]
+    ax.hist2d(delta_pos_all[:, 0], delta_pos_all[:, 1], bins=50, cmap='viridis', cmin=1)
+    ax.set_xlabel('Δx (relative position)')
+    ax.set_ylabel('Δy (relative position)')
+    ax.set_title('Relative Position Distribution')
+    ax.set_aspect('equal')
+    plt.colorbar(ax.collections[0], ax=ax, label='Count')
+
+    ax = axes[1, 2]
+    delta_vel_mag = np.linalg.norm(delta_vel_all, axis=1)
+    ax.hist(delta_vel_mag, bins=50, alpha=0.7, color='orange')
+    ax.set_xlabel('|Δv| (relative velocity magnitude)')
+    ax.set_ylabel('Frequency')
+    ax.set_title('Relative Velocity Magnitude')
+    ax.grid(True, alpha=0.3)
+    ax.axvline(delta_vel_mag.mean(), color='red', linestyle='--', label=f'Mean: {delta_vel_mag.mean():.3f}')
+    ax.legend()
+
+    # Row 3: Accelerations (targets)
+    ax = axes[2, 0]
+    acc_mag = np.linalg.norm(accelerations_all, axis=1)
+    ax.hist(acc_mag, bins=50, alpha=0.7, color='red')
+    ax.set_xlabel('Acceleration Magnitude')
+    ax.set_ylabel('Frequency')
+    ax.set_title('Acceleration Magnitude Distribution')
+    ax.grid(True, alpha=0.3)
+    ax.axvline(acc_mag.mean(), color='blue', linestyle='--', label=f'Mean: {acc_mag.mean():.4f}')
+    ax.legend()
+
+    ax = axes[2, 1]
+    ax.hist2d(accelerations_all[:, 0], accelerations_all[:, 1], bins=50, cmap='coolwarm', cmin=1)
+    ax.set_xlabel('Acceleration X')
+    ax.set_ylabel('Acceleration Y')
+    ax.set_title('Acceleration Vector Distribution')
+    ax.set_aspect('equal')
+    plt.colorbar(ax.collections[0], ax=ax, label='Count')
+
+    ax = axes[2, 2]
+    # Relative velocity vector distribution (sampled for performance)
+    sample_indices = np.random.choice(len(delta_vel_all), size=min(50000, len(delta_vel_all)), replace=False)
+    ax.hist2d(delta_vel_all[sample_indices, 0], delta_vel_all[sample_indices, 1],
+              bins=50, cmap='plasma', cmin=1)
+    ax.set_xlabel('Δv_x (relative velocity)')
+    ax.set_ylabel('Δv_y (relative velocity)')
+    ax.set_title('Relative Velocity Vector Distribution')
+    ax.set_aspect('equal')
+    plt.colorbar(ax.collections[0], ax=ax, label='Count')
+
+    plt.tight_layout()
+
+    if save_dir:
+        plot_path = os.path.join(save_dir, 'eda_distributions.png')
+        plt.savefig(plot_path, dpi=300, bbox_inches='tight')
+        logger.info(f"Saved distribution plots to {plot_path}")
+    else:
+        plt.show()
+
+    # Generate ground truth force field plots (if config is available)
+    logger.info("Generating ground truth force field plots...")
+
+    # Try to load config
+    config_path = dataset_path.replace('.pt', '_config.pt')
+    if os.path.exists(config_path):
+        logger.info(f"Found config at {config_path}")
+
+        try:
+            # Import plotting functions
+            try:
+                from .plotting import (
+                    evaluate_true_boid_forces,
+                    evaluate_true_boid_velocity_forces,
+                    plot_force_decomposition,
+                    plot_symmetric_force_decomposition
+                )
+            except ImportError:
+                # If relative import fails, try absolute import
+                from collab_env.gnn.interaction_particles.plotting import (
+                    evaluate_true_boid_forces,
+                    evaluate_true_boid_velocity_forces,
+                    plot_force_decomposition,
+                    plot_symmetric_force_decomposition
+                )
+
+            boids_config = torch.load(config_path, weights_only=False)
+            logger.info(f"Using config: {boids_config}")
+
+            # Get the first species config
+            species_key = list(boids_config.keys())[0] if boids_config else None
+            if species_key and species_key != 'scene_size':
+                logger.info(f"Using config for species '{species_key}'")
+
+                # Extract species-specific config
+                species_config = boids_config[species_key]
+
+                # Evaluate ground truth forces
+                # Determine plot boundary based on which forces are active
+                # For independent configs, use min_distance; otherwise use visual_range
+                if species_config.get('independent', False):
+                    # Independent: only avoidance active (within min_distance)
+                    relevant_range = species_config['min_distance'] / scene_size
+                else:
+                    # Normal: cohesion/alignment active (within visual_range)
+                    relevant_range = species_config['visual_range'] / scene_size
+
+                plot_max_dist = relevant_range * 1.5  # 50% cushion
+
+                true_forces = evaluate_true_boid_forces(
+                    species_config,  # Pass species-specific config
+                    grid_size=50,
+                    max_dist=plot_max_dist,
+                    scene_size=scene_size
+                )
+
+                # Plot ground truth forces
+                # Use visual_range and min_distance from config (normalized to [0,1] space)
+                visual_range_normalized = species_config['visual_range'] / scene_size
+                min_distance_normalized = None
+                if 'min_distance' in species_config:
+                    min_distance_normalized = species_config['min_distance'] / scene_size
+
+                if save_dir:
+                    true_plot_path = os.path.join(save_dir, 'true_boid_force_fields.png')
+                    plot_force_decomposition(
+                        true_forces,
+                        save_path=true_plot_path,
+                        title_prefix="Ground Truth Boid",
+                        visual_range=visual_range_normalized,
+                        min_distance=min_distance_normalized
+                    )
+                    logger.info(f"Saved force field plot to {true_plot_path}")
+                else:
+                    plot_force_decomposition(
+                        true_forces,
+                        save_path=None,
+                        title_prefix="Ground Truth Boid",
+                        visual_range=visual_range_normalized,
+                        min_distance=min_distance_normalized
+                    )
+                    plt.show()
+
+                # Generate SYMMETRIC force decomposition plots (NEW!)
+                logger.info("Generating symmetric force decomposition plots...")
+
+                # Compute typical velocity magnitude from dataset
+                # Use mean relative velocity as a reference
+                speed_stats = np.linalg.norm(velocities_all, axis=1)
+                mean_speed = speed_stats.mean()
+                logger.info(f"Mean speed from dataset: {mean_speed:.4f}")
+
+                # Estimate max relative velocity (2x mean speed as upper bound)
+                max_rel_vel = min(2.0 * mean_speed, 0.05)  # Cap at 0.05
+
+                # Position-dependent forces (with delta_vel = 0)
+                pos_forces_true = evaluate_true_boid_forces(
+                    species_config,
+                    grid_size=50,
+                    max_dist=plot_max_dist,
+                    scene_size=scene_size
+                )
+
+                # Velocity-dependent forces (with delta_pos = 0)
+                vel_forces_true = evaluate_true_boid_velocity_forces(
+                    species_config,
+                    grid_size=50,
+                    max_vel=max_rel_vel,
+                    scene_size=scene_size
+                )
+
+                # Plot symmetric decomposition
+                if save_dir:
+                    symmetric_plot_path = os.path.join(save_dir, 'true_boid_symmetric_decomposition.png')
+                    plot_symmetric_force_decomposition(
+                        pos_forces_true,
+                        vel_forces_true,
+                        save_path=symmetric_plot_path,
+                        title_prefix="Ground Truth Boid",
+                        visual_range=visual_range_normalized
+                    )
+                    logger.info(f"Saved symmetric decomposition plot to {symmetric_plot_path}")
+                else:
+                    plot_symmetric_force_decomposition(
+                        pos_forces_true,
+                        vel_forces_true,
+                        save_path=None,
+                        title_prefix="Ground Truth Boid",
+                        visual_range=visual_range_normalized
+                    )
+                    plt.show()
+
+            else:
+                logger.warning("Could not find species config in loaded config file")
+
+        except Exception as e:
+            logger.warning(f"Could not generate force field plots: {e}")
+    else:
+        logger.info(f"Config file not found at {config_path}, skipping force field plots")
+
+    logger.info("Analysis complete!")
+
+    return {
+        'positions': positions_all,
+        'velocities': velocities_all,
+        'accelerations': accelerations_all,
+        'delta_pos': delta_pos_all,
+        'distances': distances_all,
+        'delta_vel': delta_vel_all
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Exploratory Data Analysis for InteractionParticle datasets',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    python analyze_dataset.py simulated_data/boid_single_species_basic.pt
+    python analyze_dataset.py simulated_data/runpod/boid_single_species_basic.pt --save-dir analysis_runpod
+    python analyze_dataset.py simulated_data/boid_single_species_basic.pt --save-dir analysis --device mps
+        """
+    )
+
+    parser.add_argument('dataset', type=str, help='Path to dataset .pt file')
+    parser.add_argument('--save-dir', type=str, default=None,
+                       help='Directory to save analysis plots (default: show plots)')
+    parser.add_argument('--device', type=str, default='cpu',
+                       help='Device to use for computations (default: cpu)')
+    parser.add_argument('--scene-size', type=float, default=480.0,
+                       help='Scene size in pixels (for normalizing config parameters)')
+
+    args = parser.parse_args()
+
+    if not os.path.exists(args.dataset):
+        logger.error(f"Dataset not found: {args.dataset}")
+        return 1
+
+    analyze_dataset(args.dataset, save_dir=args.save_dir, device=args.device,
+                   scene_size=args.scene_size)
+
+    return 0
+
+
+if __name__ == '__main__':
+    sys.exit(main())
diff --git a/collab_env/gnn/interaction_particles/model.py b/collab_env/gnn/interaction_particles/model.py
new file mode 100644
index 0000000..d94640d
--- /dev/null
+++ b/collab_env/gnn/interaction_particles/model.py
@@ -0,0 +1,262 @@
+"""
+InteractionParticle model adapted from:
+https://github.com/saalfeldlab/decomp-gnn/blob/main/src/ParticleGraph/models/Interaction_Particle.py
+
+Model learning the acceleration of particles as a function of their relative distance and relative velocities.
+"""
+
+import torch
+import torch.nn as nn
+import torch_geometric as pyg
+import torch_geometric.utils as pyg_utils
+from torch_geometric.nn.models import MLP
+import numpy as np
+
+
+class InteractionParticle(pyg.nn.MessagePassing):
+    """
+    Interaction Network for learning particle dynamics.
+
+    Based on: https://proceedings.neurips.cc/paper/2016/hash/3147da8ab4a0437c15ef51a5cc7f2dc4-Abstract.html
+
+    This model learns pairwise interaction forces between particles using edge features.
+    The interaction function is a MLP (self.lin_edge) that maps edge features to forces.
+    Each particle has a learnable embedding (self.a).
+
+    Edge Features (NOT normalized):
+    --------------------------------
+    For an edge from particle j to particle i:
+    - delta_pos: Relative position (pos_j - pos_i) [2D vector]
+    - r: Distance ||delta_pos|| [scalar]
+    - delta_vel: Relative velocity (vel_j - vel_i) [2D vector]
+    - pos_i: Absolute position of particle i [2D vector] (NEW)
+    - embedding_i: Particle i's learnable embedding [embedding_dim vector] (only if n_particle_types > 1)
+
+    Total input size: 2 (delta_pos) + 1 (r) + 2 (delta_vel) + 2 (pos_i) + [embedding_dim if n_particle_types > 1]
+                    = 7 + [embedding_dim if n_particle_types > 1]
+
+    Output:
+    -------
+    - Force from j acting on i [2D vector]
+
+    Parameters
+    ----------
+    config : dict
+        Configuration dictionary with keys:
+        - n_particles: Number of particles
+        - n_particle_types: Number of particle types (embeddings only used if > 1)
+        - max_radius: Maximum interaction radius (not used for normalization)
+        - hidden_dim: Hidden dimension for MLP
+        - embedding_dim: Dimension of particle embeddings (ignored if n_particle_types=1)
+        - n_mp_layers: Number of MLP layers
+        - input_size: Input size for MLP (should be 7 + embedding_dim if n_particle_types > 1, else 7)
+        - output_size: Output size (2 for 2D forces)
+    device : torch.device
+        Device to run the model on
+    aggr_type : str, optional
+        Aggregation type for message passing (default: 'add')
+    dimension : int, optional
+        Spatial dimension (default: 2)
+    """
+
+    def __init__(self, config, device, aggr_type='add', dimension=2):
+        super(InteractionParticle, self).__init__(aggr=aggr_type)
+
+        self.device = device
+        # Input features: delta_pos(2) + r(1) + delta_vel(2) + pos_i(2) = 7
+        self.input_size = config.get('input_size', 7)
+        self.output_size = config.get('output_size', 2)  # 2D acceleration
+        self.hidden_dim = config.get('hidden_dim', 128)
+        self.n_layers = config.get('n_mp_layers', 3)
+        self.n_particles = config.get('n_particles', 20)
+        self.n_particle_types = config.get('n_particle_types', 1)
+        self.max_radius = config.get('max_radius', 1.0)
+        self.embedding_dim = config.get('embedding_dim', 16)
+        self.dimension = dimension
+
+        self.use_embeddings = self.embedding_dim > 0 # set to 0 to disable embeddings
+
+        # Edge interaction function (MLP)
+        # Input size includes embeddings only if n_particle_types > 1
+        mlp_input_size = self.input_size + (self.embedding_dim if self.use_embeddings else 0)
+
+        # Use PyTorch Geometric's standardized MLP implementation
+        # Benefits: Better maintained, supports various activations/norms, consistent with PyG ecosystem
+        # Create channel list: [input_size, hidden_dim, ..., hidden_dim, output_size]
+        channel_list = [mlp_input_size] + [self.hidden_dim] * (self.n_layers - 1) + [self.output_size]
+
+        self.lin_edge = MLP(
+            channel_list=channel_list,
+            act='tanh',           # Tanh activation for hidden layers
+            norm='layer_norm',    # Layer normalization
+            plain_last=True,      # No activation on output layer (linear force prediction)
+            dropout=0.0           # No dropout (can tune if overfitting)
+        ).to(self.device)
+
+        # Learnable particle embeddings (only if multiple particle types)
+        if self.use_embeddings:
+            # Shape: [n_particles, embedding_dim]
+            self.a = nn.Parameter(
+                torch.randn(self.n_particles, self.embedding_dim, device=self.device) * 0.1,
+                requires_grad=True
+            )
+        else:
+            self.a = None
+
+    def forward(self, data):
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        data : torch_geometric.data.Data
+            Graph data with attributes:
+            - x: node features [N, feature_dim]
+                Expected format: [particle_id, pos_x, pos_y, vel_x, vel_y, ...]
+            - edge_index: edge connectivity [2, E]
+
+        Returns
+        -------
+        pred : torch.Tensor
+            Predicted accelerations [N, 2]
+        """
+        x, edge_index = data.x, data.edge_index
+        edge_index, _ = pyg_utils.remove_self_loops(edge_index)
+
+        # Extract features from node features
+        # Assuming x format: [particle_id, pos_x, pos_y, vel_x, vel_y, ...]
+        particle_id = x[:, 0].long()
+        pos = x[:, 1:1+self.dimension]
+        vel = x[:, 1+self.dimension:1+2*self.dimension]
+
+        # Get particle embeddings (only if using embeddings)
+        if self.use_embeddings:
+            embedding = self.a[particle_id]
+            pred = self.propagate(edge_index, pos=pos, vel=vel, embedding=embedding)
+        else:
+            pred = self.propagate(edge_index, pos=pos, vel=vel, embedding=None)
+
+        return pred
+
+    def message(self, pos_i, pos_j, vel_i, vel_j, embedding_i):
+        """
+        Construct messages from j to i.
+
+        Edge features (NOT normalized):
+        - Relative position: delta_pos = pos_j - pos_i
+        - Relative distance: r = ||delta_pos||
+        - Relative velocity: delta_vel = vel_j - vel_i
+        - Absolute position: pos_i (NEW)
+        - Particle embedding: embedding_i (only if n_particle_types > 1)
+
+        Parameters
+        ----------
+        pos_i, pos_j : torch.Tensor
+            Positions of nodes i and j
+        vel_i, vel_j : torch.Tensor
+            Velocities of nodes i and j
+        embedding_i : torch.Tensor or None
+            Embeddings of nodes i (None if n_particle_types=1)
+
+        Returns
+        -------
+        out : torch.Tensor
+            Messages to send (force from j to i)
+        """
+        # Compute relative position and distance (NO normalization)
+        delta_pos = pos_j - pos_i
+        r = torch.sqrt(torch.sum(delta_pos ** 2, dim=1, keepdim=True))
+
+        # Compute relative velocity (NO normalization)
+        delta_vel = vel_j - vel_i
+
+        # Compute complementary boundary distances to help model learn wall bouncing
+        # Distance from right/top walls (1-x, 1-y)
+        boundary_complement = 1.0 - pos_i  # (2D)
+
+        # Construct input features
+        # Format: [delta_pos(2), r(1), delta_vel(2), pos_i(2), vel_i(2), boundary_complement(2), [embedding_i(embedding_dim)]]
+        feature_list = [
+            delta_pos,            # Relative position (2D)
+            r,                    # Distance (1D)
+            delta_vel,            # Relative velocity (2D)
+            pos_i,                # Absolute position (2D) - distance from left/bottom
+            vel_i,                # Self velocity (2D)
+            boundary_complement   # Distance from right/top (2D) - NEW: [1-x, 1-y]
+        ]
+
+        # Add embeddings only if using them
+        if self.use_embeddings and embedding_i is not None:
+            feature_list.append(embedding_i)
+
+        in_features = torch.cat(feature_list, dim=-1)
+
+        # Apply edge MLP
+        out = self.lin_edge(in_features)
+
+        return out
+
+    def update(self, aggr_out):
+        """Update node features (identity function)."""
+        return aggr_out
+
+    def evaluate_interaction(self, delta_pos, delta_vel, pos_i=None, vel_i=None, embedding_idx=0):
+        """
+        Evaluate the learned interaction function.
+
+        Parameters
+        ----------
+        delta_pos : torch.Tensor
+            Relative positions [N, 2]
+        delta_vel : torch.Tensor
+            Relative velocities [N, 2]
+        pos_i : torch.Tensor, optional
+            Absolute positions of particle i [N, 2]. If None, uses origin (0, 0)
+        vel_i : torch.Tensor, optional
+            Self velocities of particle i [N, 2]. If None, uses zero velocity
+        embedding_idx : int
+            Index of particle embedding to use (ignored if n_particle_types=1)
+
+        Returns
+        -------
+        forces : torch.Tensor
+            Interaction forces [N, 2]
+        """
+        with torch.no_grad():
+            n = delta_pos.shape[0]
+
+            # Compute distance
+            r = torch.sqrt(torch.sum(delta_pos ** 2, dim=1, keepdim=True))
+
+            # Default pos_i to origin if not provided
+            if pos_i is None:
+                pos_i = torch.zeros((n, 2), device=delta_pos.device)
+
+            # Default vel_i to zero if not provided
+            if vel_i is None:
+                vel_i = torch.zeros((n, 2), device=delta_pos.device)
+
+            # Compute complementary boundary distances
+            boundary_complement = 1.0 - pos_i  # (2D)
+
+            # Construct input features
+            feature_list = [
+                delta_pos,            # Relative position (2D)
+                r,                    # Distance (1D)
+                delta_vel,            # Relative velocity (2D)
+                pos_i,                # Absolute position (2D)
+                vel_i,                # Self velocity (2D)
+                boundary_complement   # Distance from right/top (2D)
+            ]
+
+            # Add embedding only if using embeddings
+            if self.use_embeddings:
+                embedding = self.a[embedding_idx].unsqueeze(0).repeat(n, 1)
+                feature_list.append(embedding)
+
+            in_features = torch.cat(feature_list, dim=-1)
+
+            # Evaluate MLP
+            forces = self.lin_edge(in_features)
+
+            return forces
diff --git a/collab_env/gnn/interaction_particles/plotting.py b/collab_env/gnn/interaction_particles/plotting.py
new file mode 100644
index 0000000..0a772de
--- /dev/null
+++ b/collab_env/gnn/interaction_particles/plotting.py
@@ -0,0 +1,1610 @@
+"""
+Plotting and comparison utilities for InteractionParticle model.
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+from loguru import logger
+
+# Import actual boid functions for ground truth comparison
+from collab_env.sim.boids_gnn_temp.boid import (
+    fly_towards_center,
+    avoid_others,
+    match_velocity
+)
+
+
+def boid_separation_force(distances, min_distance=15.0, avoid_factor=0.05):
+    """
+    Compute true 2D boid separation force as a function of distance.
+
+    From boid.py avoid_others():
+    - For each neighbor within min_distance:
+      force += (self_pos - other_pos) * avoid_factor
+    - This is linear in displacement, not inverse-square
+
+    Parameters
+    ----------
+    distances : np.ndarray
+        Array of distances
+    min_distance : float
+        Minimum distance threshold (default: 15 from 2D boids)
+    avoid_factor : float
+        Avoidance weight (default: 0.05 from 2D boids)
+
+    Returns
+    -------
+    forces : np.ndarray
+        Separation forces
+    """
+    forces = np.zeros_like(distances)
+    mask = distances < min_distance
+    # Linear repulsion: force = avoid_factor * distance
+    forces[mask] = avoid_factor * distances[mask]
+    return forces
+
+
+def boid_alignment_force(distances, visual_range=50.0, matching_factor=0.5):
+    """
+    Compute true 2D boid alignment force as a function of distance.
+
+    From boid.py match_velocity():
+    - For neighbors within visual_range:
+      force += (avg_velocity - self_velocity) * matching_factor
+    - This is a step function based on visual range
+
+    Note: Actual alignment depends on relative velocities.
+
+    Parameters
+    ----------
+    distances : np.ndarray
+        Array of distances
+    visual_range : float
+        Visual range for alignment (default: 50 from 2D boids)
+    matching_factor : float
+        Alignment weight (default: 0.5 from 2D boids)
+
+    Returns
+    -------
+    forces : np.ndarray
+        Alignment indicator (1 if within visual range, 0 otherwise)
+    """
+    forces = np.zeros_like(distances)
+    mask = distances < visual_range
+    forces[mask] = matching_factor
+    return forces
+
+
+def boid_cohesion_force(distances, visual_range=50.0, centering_factor=0.005):
+    """
+    Compute true 2D boid cohesion force as a function of distance.
+
+    From boid.py fly_towards_center():
+    - For neighbors within visual_range:
+      force += (center_of_mass - self_pos) * centering_factor
+    - This is linear in distance to center of mass
+
+    Note: Actual cohesion is to center of neighbors, not pairwise.
+
+    Parameters
+    ----------
+    distances : np.ndarray
+        Array of distances
+    visual_range : float
+        Visual range for cohesion (default: 50 from 2D boids)
+    centering_factor : float
+        Cohesion weight (default: 0.005 from 2D boids)
+
+    Returns
+    -------
+    forces : np.ndarray
+        Cohesion forces (proportional to distance)
+    """
+    forces = np.zeros_like(distances)
+    mask = distances < visual_range
+    # Cohesion force pulls towards center, so it increases with distance
+    forces[mask] = centering_factor * distances[mask]
+    return forces
+
+
+def plot_interaction_functions(
+    model,
+    save_path=None,
+    distances=None,
+    particle_idx=0,
+    config=None,
+    velocity_slice='stationary'
+):
+    """
+    Plot learned 2D interaction functions from the model as vector field and heatmaps.
+
+    Since boid separation/cohesion depend only on position (not velocity), we visualize
+    the interaction function for a specific velocity configuration (slice through 4D space).
+
+    Parameters
+    ----------
+    model : InteractionParticle
+        Trained model
+    save_path : str, optional
+        Path to save the plot
+    distances : np.ndarray, optional
+        Distances to evaluate (if None, uses default range)
+    particle_idx : int
+        Particle index to use for embedding
+    config : dict, optional
+        Configuration dict with boid parameters
+    velocity_slice : str
+        Which velocity configuration to visualize:
+        - 'stationary': Both particles stationary (v=0)
+        - 'aligned': Both particles moving right →
+        - 'opposing': Particles moving toward each other
+        - 'perpendicular': Velocities at right angles
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        Figure object
+    """
+    import torch
+
+    # Create a 2D grid of relative positions
+    grid_size = 50
+    max_dist = model.max_radius
+    x = np.linspace(-max_dist, max_dist, grid_size)
+    y = np.linspace(-max_dist, max_dist, grid_size)
+    X, Y = np.meshgrid(x, y)
+
+    # Flatten grid for batch processing
+    positions = np.stack([X.flatten(), Y.flatten()], axis=-1)  # [N, 2]
+    n_points = len(positions)
+
+    # Compute distances
+    distances = np.linalg.norm(positions, axis=1)
+
+    # Evaluate interaction function with specified velocity configuration
+    with torch.no_grad():
+        # Convert to tensor
+        delta_pos_t = torch.tensor(positions, dtype=torch.float32, device=model.device)
+
+        # Compute distances
+        distances_t = torch.sqrt(torch.sum(delta_pos_t ** 2, dim=1, keepdim=True))
+
+        # Set velocity configuration based on slice
+        # vel_i = 0 (particle i is stationary)
+        # vel_j depends on the slice
+        vel_j = torch.zeros((n_points, 2), device=model.device)
+
+        if velocity_slice == 'stationary':
+            # Both stationary - vel_j stays zero
+            pass
+        elif velocity_slice == 'aligned':
+            # vel_j moving right
+            vel_j[:, 0] = 1.0
+        elif velocity_slice == 'opposing':
+            # vel_j pointing toward i (opposite of position vector)
+            vel_j = -delta_pos_t / (distances_t + 1e-8)
+        elif velocity_slice == 'perpendicular':
+            # vel_j perpendicular to position vector
+            vel_j[:, 0] = -delta_pos_t[:, 1]
+            vel_j[:, 1] = delta_pos_t[:, 0]
+            vel_j = vel_j / (torch.norm(vel_j, dim=1, keepdim=True) + 1e-8)
+
+        # Compute delta_vel = vel_j - vel_i (vel_i = 0, so delta_vel = vel_j)
+        delta_vel = vel_j
+
+        # Use model's evaluate_interaction method
+        forces = model.evaluate_interaction(delta_pos_t, delta_vel, embedding_idx=particle_idx).cpu().numpy()
+
+    # Reshape forces back to grid
+    force_x = forces[:, 0].reshape(grid_size, grid_size)
+    force_y = forces[:, 1].reshape(grid_size, grid_size)
+    force_magnitude = np.sqrt(force_x**2 + force_y**2)
+
+    # Compute radial and tangential components
+    # Radial: component in direction of delta_pos
+    # Tangential: component perpendicular to delta_pos
+    radial_force = np.zeros_like(force_magnitude)
+    tangential_force = np.zeros_like(force_magnitude)
+
+    for i in range(grid_size):
+        for j in range(grid_size):
+            dx, dy = X[i, j], Y[i, j]
+            dist = np.sqrt(dx**2 + dy**2)
+            if dist > 1e-8:
+                # Unit vector in radial direction
+                r_hat = np.array([dx, dy]) / dist
+                # Unit vector in tangential direction (perpendicular)
+                t_hat = np.array([-dy, dx]) / dist
+
+                force_vec = np.array([force_x[i, j], force_y[i, j]])
+                radial_force[i, j] = np.dot(force_vec, r_hat)
+                tangential_force[i, j] = np.dot(force_vec, t_hat)
+
+    # Create figure with 6 subplots (2 rows x 3 cols)
+    fig = plt.figure(figsize=(18, 11))
+
+    # Velocity mode description
+    vel_desc = {
+        'stationary': 'Stationary particles (v≈0)',
+        'aligned': 'Both particles moving right →',
+        'opposing': 'Particles moving toward each other ← →',
+        'perpendicular': 'Perpendicular velocities (→ and ↑)'
+    }
+
+    # 1. Vector field plot
+    ax1 = plt.subplot(2, 3, 1)
+    skip = 3
+    quiver = ax1.quiver(X[::skip, ::skip], Y[::skip, ::skip],
+                        force_x[::skip, ::skip], force_y[::skip, ::skip],
+                        force_magnitude[::skip, ::skip],
+                        cmap='viridis', scale=None, scale_units='xy')
+    ax1.set_xlabel('Relative Position X', fontsize=10)
+    ax1.set_ylabel('Relative Position Y', fontsize=10)
+    ax1.set_title('Force Vector Field', fontsize=11, fontweight='bold')
+    ax1.set_aspect('equal')
+    ax1.grid(True, alpha=0.3)
+    plt.colorbar(quiver, ax=ax1, label='Magnitude')
+
+    # 2. Force magnitude heatmap
+    ax2 = plt.subplot(2, 3, 2)
+    im2 = ax2.contourf(X, Y, force_magnitude, levels=20, cmap='hot')
+    ax2.set_xlabel('Relative Position X', fontsize=10)
+    ax2.set_ylabel('Relative Position Y', fontsize=10)
+    ax2.set_title('Force Magnitude', fontsize=11, fontweight='bold')
+    ax2.set_aspect('equal')
+    plt.colorbar(im2, ax=ax2, label='Magnitude')
+
+    # 3. Radial component (repulsion/attraction)
+    ax3 = plt.subplot(2, 3, 3)
+    vmax = np.max(np.abs(radial_force))
+    im3 = ax3.contourf(X, Y, radial_force, levels=20, cmap='RdBu_r',
+                       vmin=-vmax, vmax=vmax)
+    ax3.set_xlabel('Relative Position X', fontsize=10)
+    ax3.set_ylabel('Relative Position Y', fontsize=10)
+    ax3.set_title('Radial Component\n(+: repulsion, -: attraction)', fontsize=11, fontweight='bold')
+    ax3.set_aspect('equal')
+    plt.colorbar(im3, ax=ax3, label='Force')
+
+    # 4. Force X component heatmap
+    ax4 = plt.subplot(2, 3, 4)
+    im4 = ax4.contourf(X, Y, force_x, levels=20, cmap='RdBu_r',
+                       vmin=-np.max(np.abs(force_x)), vmax=np.max(np.abs(force_x)))
+    ax4.set_xlabel('Relative Position X', fontsize=10)
+    ax4.set_ylabel('Relative Position Y', fontsize=10)
+    ax4.set_title('Force X Component', fontsize=11, fontweight='bold')
+    ax4.set_aspect('equal')
+    plt.colorbar(im4, ax=ax4, label='Force X')
+
+    # 5. Force Y component heatmap
+    ax5 = plt.subplot(2, 3, 5)
+    im5 = ax5.contourf(X, Y, force_y, levels=20, cmap='RdBu_r',
+                       vmin=-np.max(np.abs(force_y)), vmax=np.max(np.abs(force_y)))
+    ax5.set_xlabel('Relative Position X', fontsize=10)
+    ax5.set_ylabel('Relative Position Y', fontsize=10)
+    ax5.set_title('Force Y Component', fontsize=11, fontweight='bold')
+    ax5.set_aspect('equal')
+    plt.colorbar(im5, ax=ax5, label='Force Y')
+
+    # 6. Tangential component
+    ax6 = plt.subplot(2, 3, 6)
+    vmax_tan = np.max(np.abs(tangential_force))
+    im6 = ax6.contourf(X, Y, tangential_force, levels=20, cmap='RdBu_r',
+                       vmin=-vmax_tan, vmax=vmax_tan)
+    ax6.set_xlabel('Relative Position X', fontsize=10)
+    ax6.set_ylabel('Relative Position Y', fontsize=10)
+    ax6.set_title('Tangential Component\n(perpendicular to radial)', fontsize=11, fontweight='bold')
+    ax6.set_aspect('equal')
+    plt.colorbar(im6, ax=ax6, label='Force')
+
+    plt.suptitle(f'Learned 2D Interaction Function (Particle {particle_idx})\nVelocity slice: {vel_desc.get(velocity_slice, velocity_slice)}',
+                 fontsize=13, fontweight='bold', y=0.995)
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        logger.info(f"Saved interaction function plot to {save_path}")
+
+    return fig
+
+
+def plot_true_boid_rules_2d(config=None, scene_size=480.0, save_path=None):
+    """
+    Visualize true 2D boid rules as 2D heatmaps.
+
+    Separation and cohesion depend only on position (radially symmetric).
+    Alignment depends on relative velocity (not visualized here as it's velocity-dependent).
+
+    Parameters
+    ----------
+    config : dict, optional
+        Boid configuration dict
+    scene_size : float
+        Scene size in pixels
+    save_path : str, optional
+        Path to save plot
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        Figure object
+    """
+    # Default config
+    if config is None:
+        config = {
+            'visual_range': 50.0,
+            'min_distance': 15.0,
+            'avoid_factor': 0.05,
+            'matching_factor': 0.5,
+            'centering_factor': 0.005,
+        }
+
+    # Create 2D grid in pixel space
+    grid_size = 100
+    visual_range_normalized = config['visual_range'] / scene_size
+    x = np.linspace(-visual_range_normalized, visual_range_normalized, grid_size)
+    y = np.linspace(-visual_range_normalized, visual_range_normalized, grid_size)
+    X, Y = np.meshgrid(x, y)
+
+    # Compute distances in pixels
+    distances_px = np.sqrt(X**2 + Y**2) * scene_size
+
+    # Compute true boid forces
+    separation = np.zeros_like(distances_px)
+    cohesion = np.zeros_like(distances_px)
+
+    # Separation (repulsion within min_distance)
+    mask_sep = distances_px < config['min_distance']
+    separation[mask_sep] = config['avoid_factor'] * distances_px[mask_sep]
+
+    # Cohesion (attraction within visual_range)
+    mask_coh = distances_px < config['visual_range']
+    cohesion[mask_coh] = config['centering_factor'] * distances_px[mask_coh]
+
+    # Convert to force vectors (radial direction)
+    # Separation: pushes away (positive radial)
+    # Cohesion: pulls toward (negative radial)
+    sep_force_x = np.zeros_like(X)
+    sep_force_y = np.zeros_like(Y)
+    coh_force_x = np.zeros_like(X)
+    coh_force_y = np.zeros_like(Y)
+
+    for i in range(grid_size):
+        for j in range(grid_size):
+            dist = np.sqrt(X[i,j]**2 + Y[i,j]**2)
+            if dist > 1e-8:
+                # Unit vector in radial direction
+                r_hat_x = X[i,j] / dist
+                r_hat_y = Y[i,j] / dist
+
+                # Separation pushes away
+                sep_force_x[i,j] = separation[i,j] * r_hat_x * scene_size
+                sep_force_y[i,j] = separation[i,j] * r_hat_y * scene_size
+
+                # Cohesion pulls toward (negative)
+                coh_force_x[i,j] = -cohesion[i,j] * r_hat_x * scene_size
+                coh_force_y[i,j] = -cohesion[i,j] * r_hat_y * scene_size
+
+    # Combined force
+    total_force_x = sep_force_x + coh_force_x
+    total_force_y = sep_force_y + coh_force_y
+    total_magnitude = np.sqrt(total_force_x**2 + total_force_y**2)
+    sep_magnitude = np.sqrt(sep_force_x**2 + sep_force_y**2)
+    coh_magnitude = np.sqrt(coh_force_x**2 + coh_force_y**2)
+
+    # Create figure
+    fig = plt.figure(figsize=(18, 12))
+
+    # 1. Separation force
+    ax1 = plt.subplot(2, 3, 1)
+    skip = 5
+    quiver1 = ax1.quiver(X[::skip, ::skip]*scene_size, Y[::skip, ::skip]*scene_size,
+                         sep_force_x[::skip, ::skip], sep_force_y[::skip, ::skip],
+                         sep_magnitude[::skip, ::skip],
+                         cmap='Reds', scale=None, scale_units='xy')
+    circle1 = plt.Circle((0, 0), config['min_distance'], fill=False,
+                         edgecolor='red', linestyle='--', linewidth=2, label='min_distance')
+    ax1.add_patch(circle1)
+    ax1.set_xlabel('Relative Position X (px)', fontsize=10)
+    ax1.set_ylabel('Relative Position Y (px)', fontsize=10)
+    ax1.set_title('Separation Force\n(Repulsion within min_distance)', fontsize=11, fontweight='bold')
+    ax1.set_aspect('equal')
+    ax1.grid(True, alpha=0.3)
+    ax1.legend()
+    plt.colorbar(quiver1, ax=ax1, label='Force')
+
+    # 2. Separation magnitude heatmap
+    ax2 = plt.subplot(2, 3, 2)
+    im2 = ax2.contourf(X*scene_size, Y*scene_size, sep_magnitude, levels=20, cmap='Reds')
+    circle2 = plt.Circle((0, 0), config['min_distance'], fill=False,
+                         edgecolor='darkred', linestyle='--', linewidth=2)
+    ax2.add_patch(circle2)
+    ax2.set_xlabel('Relative Position X (px)', fontsize=10)
+    ax2.set_ylabel('Relative Position Y (px)', fontsize=10)
+    ax2.set_title('Separation Magnitude', fontsize=11, fontweight='bold')
+    ax2.set_aspect('equal')
+    plt.colorbar(im2, ax=ax2, label='Force')
+
+    # 3. Cohesion force
+    ax3 = plt.subplot(2, 3, 3)
+    quiver3 = ax3.quiver(X[::skip, ::skip]*scene_size, Y[::skip, ::skip]*scene_size,
+                         coh_force_x[::skip, ::skip], coh_force_y[::skip, ::skip],
+                         coh_magnitude[::skip, ::skip],
+                         cmap='Blues_r', scale=None, scale_units='xy')
+    circle3 = plt.Circle((0, 0), config['visual_range'], fill=False,
+                         edgecolor='blue', linestyle='--', linewidth=2, label='visual_range')
+    ax3.add_patch(circle3)
+    ax3.set_xlabel('Relative Position X (px)', fontsize=10)
+    ax3.set_ylabel('Relative Position Y (px)', fontsize=10)
+    ax3.set_title('Cohesion Force\n(Attraction within visual_range)', fontsize=11, fontweight='bold')
+    ax3.set_aspect('equal')
+    ax3.grid(True, alpha=0.3)
+    ax3.legend()
+    plt.colorbar(quiver3, ax=ax3, label='Force')
+
+    # 4. Cohesion magnitude heatmap
+    ax4 = plt.subplot(2, 3, 4)
+    im4 = ax4.contourf(X*scene_size, Y*scene_size, coh_magnitude, levels=20, cmap='Blues_r')
+    circle4 = plt.Circle((0, 0), config['visual_range'], fill=False,
+                         edgecolor='darkblue', linestyle='--', linewidth=2)
+    ax4.add_patch(circle4)
+    ax4.set_xlabel('Relative Position X (px)', fontsize=10)
+    ax4.set_ylabel('Relative Position Y (px)', fontsize=10)
+    ax4.set_title('Cohesion Magnitude', fontsize=11, fontweight='bold')
+    ax4.set_aspect('equal')
+    plt.colorbar(im4, ax=ax4, label='Force')
+
+    # 5. Combined force
+    ax5 = plt.subplot(2, 3, 5)
+    quiver5 = ax5.quiver(X[::skip, ::skip]*scene_size, Y[::skip, ::skip]*scene_size,
+                         total_force_x[::skip, ::skip], total_force_y[::skip, ::skip],
+                         total_magnitude[::skip, ::skip],
+                         cmap='viridis', scale=None, scale_units='xy')
+    circle5a = plt.Circle((0, 0), config['min_distance'], fill=False,
+                          edgecolor='red', linestyle='--', linewidth=1.5, alpha=0.7)
+    circle5b = plt.Circle((0, 0), config['visual_range'], fill=False,
+                          edgecolor='blue', linestyle='--', linewidth=1.5, alpha=0.7)
+    ax5.add_patch(circle5a)
+    ax5.add_patch(circle5b)
+    ax5.set_xlabel('Relative Position X (px)', fontsize=10)
+    ax5.set_ylabel('Relative Position Y (px)', fontsize=10)
+    ax5.set_title('Combined Force\n(Separation + Cohesion)', fontsize=11, fontweight='bold')
+    ax5.set_aspect('equal')
+    ax5.grid(True, alpha=0.3)
+    plt.colorbar(quiver5, ax=ax5, label='Force')
+
+    # 6. Combined magnitude heatmap
+    ax6 = plt.subplot(2, 3, 6)
+    im6 = ax6.contourf(X*scene_size, Y*scene_size, total_magnitude, levels=20, cmap='viridis')
+    circle6a = plt.Circle((0, 0), config['min_distance'], fill=False,
+                          edgecolor='red', linestyle='--', linewidth=1.5, alpha=0.7)
+    circle6b = plt.Circle((0, 0), config['visual_range'], fill=False,
+                          edgecolor='blue', linestyle='--', linewidth=1.5, alpha=0.7)
+    ax6.add_patch(circle6a)
+    ax6.add_patch(circle6b)
+    ax6.set_xlabel('Relative Position X (px)', fontsize=10)
+    ax6.set_ylabel('Relative Position Y (px)', fontsize=10)
+    ax6.set_title('Combined Magnitude', fontsize=11, fontweight='bold')
+    ax6.set_aspect('equal')
+    plt.colorbar(im6, ax=ax6, label='Force')
+
+    plt.suptitle('True 2D Boid Rules (Position-Dependent Forces)',
+                 fontsize=14, fontweight='bold', y=0.995)
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        logger.info(f"Saved true boid rules plot to {save_path}")
+
+    return fig
+
+
+def evaluate_forces_on_grid(model, grid_size=60, max_dist=50.0, particle_idx=0, speed_magnitude=0.015):
+    """
+    Evaluate learned forces on a 2D grid for two velocity configurations.
+
+    Setup:
+    - Node i at square center (0.5, 0.5), stationary (vel_i = 0)
+    - Node j at grid positions relative to i (delta_pos varies)
+
+    Parameters
+    ----------
+    model : InteractionParticle
+        Trained model
+    grid_size : int
+        Grid resolution
+    max_dist : float
+        Maximum distance to evaluate (in same units as training data)
+    particle_idx : int
+        Particle embedding index
+    speed_magnitude : float
+        Magnitude of relative velocity to use (should match typical training data speed)
+        Default: 0.015 (typical for normalized boid data)
+
+    Returns
+    -------
+    dict with keys:
+        - X, Y: meshgrid coordinates
+        - away_total, away_pos, away_vel: forces when j moves away from i
+        - towards_total, towards_pos, towards_vel: forces when j moves towards i
+    """
+    # Create 2D grid
+    x = np.linspace(-max_dist, max_dist, grid_size)
+    y = np.linspace(-max_dist, max_dist, grid_size)
+    X, Y = np.meshgrid(x, y)
+
+    # Flatten for batch processing
+    positions = np.stack([X.flatten(), Y.flatten()], axis=-1)  # [N, 2]
+    n_points = len(positions)
+
+    # Convert to tensors
+    delta_pos_t = torch.tensor(positions, dtype=torch.float32, device=model.device)
+
+    # Set absolute position to square center (0.5, 0.5)
+    # This is important because the model uses absolute position as a feature
+    pos_i = torch.full((n_points, 2), 0.5, dtype=torch.float32, device=model.device)
+
+    # Set particle i's velocity to zero (stationary)
+    # This is important because the model uses self velocity as a feature
+    vel_i = torch.zeros((n_points, 2), dtype=torch.float32, device=model.device)
+
+    # Compute distances and velocity directions with REALISTIC MAGNITUDES
+    distances = torch.sqrt(torch.sum(delta_pos_t ** 2, dim=1, keepdim=True))
+    r_hat = delta_pos_t / (distances + 1e-8)  # Unit vector away from i
+
+    # Use realistic speed magnitude (not unit vectors!)
+    vel_away = r_hat * speed_magnitude  # Realistic velocity away from i
+    vel_towards = -r_hat * speed_magnitude  # Realistic velocity towards i
+    vel_zero = torch.zeros_like(delta_pos_t)
+
+    # Evaluate forces for all configurations
+    # NOTE: Model computes force ON i (at square center) DUE TO j (at grid position)
+    # But we want to plot force ON j DUE TO i (Newton's 3rd law: F_j = -F_i)
+    # So we negate the model output to match the true boid force convention
+
+    # 1. Away: j moving away from i
+    away_total = -model.evaluate_interaction(delta_pos_t, vel_away, pos_i=pos_i, vel_i=vel_i, embedding_idx=particle_idx).cpu().numpy()
+    away_pos = -model.evaluate_interaction(delta_pos_t, vel_zero, pos_i=pos_i, vel_i=vel_i, embedding_idx=particle_idx).cpu().numpy()
+    away_vel = away_total - away_pos
+
+    # 2. Towards: j moving towards i
+    towards_total = -model.evaluate_interaction(delta_pos_t, vel_towards, pos_i=pos_i, vel_i=vel_i, embedding_idx=particle_idx).cpu().numpy()
+    towards_pos = -model.evaluate_interaction(delta_pos_t, vel_zero, pos_i=pos_i, vel_i=vel_i, embedding_idx=particle_idx).cpu().numpy()
+    towards_vel = towards_total - towards_pos
+
+    # Reshape to grid
+    def reshape_to_grid(arr):
+        return {
+            'x': arr[:, 0].reshape(grid_size, grid_size),
+            'y': arr[:, 1].reshape(grid_size, grid_size),
+            'mag': np.sqrt(arr[:, 0]**2 + arr[:, 1]**2).reshape(grid_size, grid_size)
+        }
+
+    return {
+        'X': X,
+        'Y': Y,
+        'away_total': reshape_to_grid(away_total),
+        'away_pos': reshape_to_grid(away_pos),
+        'away_vel': reshape_to_grid(away_vel),
+        'towards_total': reshape_to_grid(towards_total),
+        'towards_pos': reshape_to_grid(towards_pos),
+        'towards_vel': reshape_to_grid(towards_vel),
+    }
+
+
+def evaluate_velocity_forces_on_grid(model, grid_size=60, max_vel=0.05, particle_idx=0):
+    """
+    Evaluate learned forces on a 2D velocity grid (symmetric decomposition).
+
+    Setup:
+    - Particles at SAME position (delta_pos = 0) at square center (0.5, 0.5)
+    - Vary relative velocity on a 2D grid
+
+    Parameters
+    ----------
+    model : InteractionParticle
+        Trained model
+    grid_size : int
+        Grid resolution
+    max_vel : float
+        Maximum relative velocity magnitude to evaluate
+    particle_idx : int
+        Particle embedding index
+
+    Returns
+    -------
+    dict with keys:
+        - VX, VY: meshgrid coordinates (relative velocity components)
+        - forces: force vectors [grid_size, grid_size, 2]
+        - force_mag: force magnitudes [grid_size, grid_size]
+    """
+    # Create 2D velocity grid
+    vx = np.linspace(-max_vel, max_vel, grid_size)
+    vy = np.linspace(-max_vel, max_vel, grid_size)
+    VX, VY = np.meshgrid(vx, vy)
+
+    # Flatten for batch processing
+    delta_vels = np.stack([VX.flatten(), VY.flatten()], axis=-1)  # [N, 2]
+    n_points = len(delta_vels)
+
+    # Convert to tensors
+    delta_vel_t = torch.tensor(delta_vels, dtype=torch.float32, device=model.device)
+
+    # Particles at same position: delta_pos = 0
+    delta_pos_t = torch.zeros_like(delta_vel_t)
+
+    # Set absolute position to square center (0.5, 0.5)
+    pos_i = torch.full((n_points, 2), 0.5, dtype=torch.float32, device=model.device)
+
+    # Set particle i's velocity to zero (stationary)
+    vel_i = torch.zeros((n_points, 2), dtype=torch.float32, device=model.device)
+
+    # Evaluate forces
+    # NOTE: Model computes force ON i DUE TO j
+    # With delta_pos=0, both particles are at same location (square center)
+    # Force should depend only on relative velocity
+    forces = -model.evaluate_interaction(delta_pos_t, delta_vel_t, pos_i=pos_i, vel_i=vel_i, embedding_idx=particle_idx).cpu().numpy()
+
+    # Reshape to grid
+    forces_x = forces[:, 0].reshape(grid_size, grid_size)
+    forces_y = forces[:, 1].reshape(grid_size, grid_size)
+    forces_mag = np.sqrt(forces_x**2 + forces_y**2)
+
+    return {
+        'VX': VX,
+        'VY': VY,
+        'forces': forces.reshape(grid_size, grid_size, 2),
+        'forces_x': forces_x,
+        'forces_y': forces_y,
+        'force_mag': forces_mag
+    }
+
+
+def plot_force_decomposition(forces_dict, save_path=None, title_prefix="Learned", visual_range=None, min_distance=None):
+    """
+    Plot force decomposition for both velocity slices.
+
+    Creates a 2x3 grid:
+    - Row 1: Away scenario (Total | Position-Only | Velocity Residual)
+    - Row 2: Towards scenario (Total | Position-Only | Velocity Residual)
+
+    Parameters
+    ----------
+    forces_dict : dict
+        Output from evaluate_forces_on_grid or evaluate_true_boid_forces
+    save_path : str, optional
+        Path to save figure
+    title_prefix : str
+        Prefix for title (e.g., "Learned" or "True Boid")
+    visual_range : float, optional
+        Visual range radius to draw as reference circle (in normalized units)
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+    """
+    X = forces_dict['X']
+    Y = forces_dict['Y']
+
+    fig = plt.figure(figsize=(18, 12))
+
+    # Common settings
+    skip = 3  # For quiver plots
+
+    def plot_force_field(ax, force_dict, title):
+        """Helper to plot a single force field."""
+        fx, fy, mag = force_dict['x'], force_dict['y'], force_dict['mag']
+
+        # Filter out near-zero vectors to help with autoscaling
+        # Define threshold as 1% of max magnitude
+        max_mag = np.max(mag)
+        if max_mag > 1e-10:
+            threshold = max_mag * 0.01
+        else:
+            threshold = 1e-10
+
+        # Apply skip and create mask for non-zero vectors
+        X_skip = X[::skip, ::skip]
+        Y_skip = Y[::skip, ::skip]
+        fx_skip = fx[::skip, ::skip]
+        fy_skip = fy[::skip, ::skip]
+        mag_skip = mag[::skip, ::skip]
+
+        # Mask for vectors above threshold
+        mask = mag_skip > threshold
+
+        # Only plot non-trivial vectors
+        if np.sum(mask) > 0:
+            X_plot = X_skip[mask]
+            Y_plot = Y_skip[mask]
+            fx_plot = fx_skip[mask]
+            fy_plot = fy_skip[mask]
+            mag_plot = mag_skip[mask]
+
+            # Vector field with autoscaling (filtering fixed the autoscale issue)
+            quiver = ax.quiver(X_plot, Y_plot, fx_plot, fy_plot, mag_plot,
+                              cmap='viridis', scale=None, scale_units='xy', alpha=0.8)
+        else:
+            # Fallback: create empty quiver
+            quiver = ax.quiver([], [], [], [], [], cmap='viridis')
+
+        # Mark origin (node i position)
+        ax.scatter([0], [0], c='red', s=200, marker='o', edgecolors='black',
+                  linewidth=2, zorder=10, label='Node i (origin)')
+
+        # Draw visual range circle if provided
+        if visual_range is not None:
+            circle = plt.Circle((0, 0), visual_range, fill=False,
+                              edgecolor='gray', linestyle='--', linewidth=1.5,
+                              alpha=0.6, zorder=5, label=f'Visual range ({visual_range:.3f})')
+            ax.add_patch(circle)
+
+        # Draw min_distance circle if provided (for true boid forces)
+        if min_distance is not None:
+            min_circle = plt.Circle((0, 0), min_distance, fill=False,
+                                   edgecolor='blue', linestyle='--', linewidth=1.5,
+                                   alpha=0.6, zorder=5, label=f'Min distance ({min_distance:.3f})')
+            ax.add_patch(min_circle)
+
+        ax.set_xlabel('Relative x position (j - i)', fontsize=10)
+        ax.set_ylabel('Relative y position (j - i)', fontsize=10)
+        ax.set_title(title, fontsize=11, fontweight='bold')
+        ax.set_aspect('equal')
+        ax.grid(True, alpha=0.3)
+        ax.legend(loc='upper right', fontsize=8)
+
+        return quiver
+
+    # Row 1: Away scenario
+    ax1 = plt.subplot(2, 3, 1)
+    q1 = plot_force_field(ax1, forces_dict['away_total'], 'I) Total Force\n(vel_j away from i)')
+    plt.colorbar(q1, ax=ax1, label='Force mag')
+
+    ax2 = plt.subplot(2, 3, 2)
+    q2 = plot_force_field(ax2, forces_dict['away_pos'], 'II) Position-Only\n(vel=0)')
+    plt.colorbar(q2, ax=ax2, label='Force mag')
+
+    ax3 = plt.subplot(2, 3, 3)
+    q3 = plot_force_field(ax3, forces_dict['away_vel'], 'III) Velocity Residual\n(I - II)')
+    plt.colorbar(q3, ax=ax3, label='Force mag')
+
+    # Row 2: Towards scenario
+    ax4 = plt.subplot(2, 3, 4)
+    q4 = plot_force_field(ax4, forces_dict['towards_total'], 'I) Total Force\n(vel_j towards i)')
+    plt.colorbar(q4, ax=ax4, label='Force mag')
+
+    ax5 = plt.subplot(2, 3, 5)
+    q5 = plot_force_field(ax5, forces_dict['towards_pos'], 'II) Position-Only\n(vel=0)')
+    plt.colorbar(q5, ax=ax5, label='Force mag')
+
+    ax6 = plt.subplot(2, 3, 6)
+    q6 = plot_force_field(ax6, forces_dict['towards_vel'], 'III) Velocity Residual\n(I - II)')
+    plt.colorbar(q6, ax=ax6, label='Force mag')
+
+    plt.suptitle(f'{title_prefix} Force on j due to i at origin (arrows show force on j at that position)',
+                 fontsize=14, fontweight='bold', y=0.995)
+    plt.tight_layout(rect=[0, 0, 1, 0.99])
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        logger.info(f"Saved force decomposition plot to {save_path}")
+
+    return fig
+
+
+def evaluate_true_boid_forces(config, grid_size=60, max_dist=50.0, scene_size=480.0):
+    """
+    Evaluate true 2D boid forces on a grid using actual boid.py functions.
+
+    Computes the force experienced by particle j (at grid positions) due to
+    particle i (at origin). This matches the natural interpretation: arrows show
+    the force on a particle at that position relative to a reference particle at origin.
+
+    Computes forces from the 3 boid rules:
+    1. Cohesion (fly_towards_center): attraction within visual_range
+    2. Separation (avoid_others): repulsion within min_distance
+    3. Alignment (match_velocity): velocity matching within visual_range
+
+    Parameters
+    ----------
+    config : dict
+        Boid config with parameters in PIXEL SPACE:
+        - min_distance: Separation threshold (pixels)
+        - visual_range: Cohesion and alignment threshold (pixels)
+        - avoid_factor: Separation weight
+        - centering_factor: Cohesion weight
+        - matching_factor: Alignment weight
+        - speed_limit: Maximum speed (pixels/timestep)
+        - independent: Whether to skip cohesion/alignment (default: False)
+    grid_size : int
+        Grid resolution
+    max_dist : float
+        Maximum distance in normalized space (recommend: visual_range * 1.5 / scene_size)
+    scene_size : float
+        Scene size in pixels (used to normalize config parameters)
+
+    Returns
+    -------
+    dict with same structure as evaluate_forces_on_grid
+    """
+    # Create 2D grid in normalized space
+    x = np.linspace(-max_dist, max_dist, grid_size)
+    y = np.linspace(-max_dist, max_dist, grid_size)
+    X, Y = np.meshgrid(x, y)
+
+    # Normalize config parameters from pixel space to normalized space
+    # The grid is in normalized space [-max_dist, max_dist], so we need to
+    # convert pixel-space thresholds to normalized space
+    normalized_config = config.copy()
+    if 'min_distance' in config:
+        normalized_config['min_distance'] = config['min_distance'] / scene_size
+    if 'visual_range' in config:
+        normalized_config['visual_range'] = config['visual_range'] / scene_size
+
+    # Flatten for iteration
+    positions = np.stack([X.flatten(), Y.flatten()], axis=-1)
+    n_points = len(positions)
+
+    # Compute radial unit vectors (direction from i to j)
+    distances = np.sqrt(X**2 + Y**2).flatten()
+    r_hat = np.zeros_like(positions)
+    mask_nonzero = distances > 1e-8
+    r_hat[mask_nonzero] = positions[mask_nonzero] / distances[mask_nonzero, np.newaxis]
+
+    # Add 'independent' flag if not present in normalized config
+    if 'independent' not in normalized_config:
+        normalized_config['independent'] = False
+
+    # Function to compute forces using actual boid functions
+    def compute_boid_forces(vel_j_array):
+        """
+        Compute forces for all grid points using actual boid.py functions.
+
+        vel_j_array: [n_points, 2] array of velocities for particle j
+        """
+        forces_total = np.zeros_like(positions)
+        forces_pos_only = np.zeros_like(positions)
+
+        for idx in range(n_points):
+            # Create boid i at origin (stationary reference point)
+            boid_i = {
+                'x': 0.0,
+                'y': 0.0,
+                'dx': 0.0,
+                'dy': 0.0,
+                'species': 'A'
+            }
+
+            # Create boid j at grid position (this is the boid we compute forces FOR)
+            pos_j = positions[idx]
+            vel_j = vel_j_array[idx]
+
+            boid_j = {
+                'x': pos_j[0],
+                'y': pos_j[1],
+                'dx': vel_j[0],
+                'dy': vel_j[1],
+                'species': 'A'
+            }
+
+            # Compute force ON boid_j (at grid position) due TO boid_i (at origin)
+            # Put boid_j in the list and call force functions on it
+            boids = [boid_j, boid_i]
+
+            fly_towards_center(boid_j, boids, normalized_config)  # Rule 1: Cohesion
+            avoid_others(boid_j, boids, normalized_config)        # Rule 2: Separation
+            match_velocity(boid_j, boids, normalized_config)      # Rule 3: Alignment
+
+            # Force is change in velocity (from initial velocity)
+            forces_total[idx] = [boid_j['dx'] - vel_j[0], boid_j['dy'] - vel_j[1]]
+
+            # Compute position-only force (set velocities to zero)
+            boid_j_zero = {
+                'x': pos_j[0],
+                'y': pos_j[1],
+                'dx': 0.0,
+                'dy': 0.0,
+                'species': 'A'
+            }
+            boid_i_zero = {
+                'x': 0.0,
+                'y': 0.0,
+                'dx': 0.0,
+                'dy': 0.0,
+                'species': 'A'
+            }
+            boids_zero = [boid_j_zero, boid_i_zero]
+
+            fly_towards_center(boid_j_zero, boids_zero, normalized_config)
+            avoid_others(boid_j_zero, boids_zero, normalized_config)
+            match_velocity(boid_j_zero, boids_zero, normalized_config)
+
+            forces_pos_only[idx] = [boid_j_zero['dx'], boid_j_zero['dy']]
+
+        return forces_total, forces_pos_only
+
+    # Use realistic velocity magnitude from config's speed_limit instead of unit vectors
+    # The speed_limit is in pixel space, so normalize it
+    if 'speed_limit' in config:
+        speed_magnitude = config['speed_limit'] / scene_size
+    else:
+        # Fallback: use a typical speed (mean from basic datasets ~0.014)
+        speed_magnitude = 0.015
+
+    # Scenario 1: vel_j moving away from i (with realistic speed)
+    vel_away = r_hat.copy() * speed_magnitude
+    forces_away_total, forces_away_pos = compute_boid_forces(vel_away)
+    forces_away_vel = forces_away_total - forces_away_pos
+
+    # Scenario 2: vel_j moving towards i (with realistic speed)
+    vel_towards = -r_hat.copy() * speed_magnitude
+    forces_towards_total, forces_towards_pos = compute_boid_forces(vel_towards)
+    forces_towards_vel = forces_towards_total - forces_towards_pos
+
+    # Reshape to grid
+    def make_force_dict(forces):
+        return {
+            'x': forces[:, 0].reshape(grid_size, grid_size),
+            'y': forces[:, 1].reshape(grid_size, grid_size),
+            'mag': np.linalg.norm(forces, axis=1).reshape(grid_size, grid_size)
+        }
+
+    return {
+        'X': X,
+        'Y': Y,
+        'away_total': make_force_dict(forces_away_total),
+        'away_pos': make_force_dict(forces_away_pos),
+        'away_vel': make_force_dict(forces_away_vel),
+        'towards_total': make_force_dict(forces_towards_total),
+        'towards_pos': make_force_dict(forces_towards_pos),
+        'towards_vel': make_force_dict(forces_towards_vel),
+    }
+
+
+def evaluate_true_boid_velocity_forces(config, grid_size=60, max_vel=0.05, scene_size=480.0):
+    """
+    Evaluate true boid forces on a 2D velocity grid (symmetric decomposition).
+
+    Setup:
+    - Two boids at SAME position (delta_pos = 0) at square center
+    - Vary relative velocity on a 2D grid
+
+    Parameters
+    ----------
+    config : dict
+        Boid config with parameters in PIXEL SPACE (same as evaluate_true_boid_forces)
+    grid_size : int
+        Grid resolution
+    max_vel : float
+        Maximum relative velocity magnitude to evaluate (in normalized space)
+    scene_size : float
+        Scene size in pixels (used to normalize config parameters)
+
+    Returns
+    -------
+    dict with keys:
+        - VX, VY: meshgrid coordinates (relative velocity components)
+        - forces: force vectors [grid_size, grid_size, 2]
+        - force_mag: force magnitudes [grid_size, grid_size]
+    """
+    # Create 2D velocity grid in normalized space
+    vx = np.linspace(-max_vel, max_vel, grid_size)
+    vy = np.linspace(-max_vel, max_vel, grid_size)
+    VX, VY = np.meshgrid(vx, vy)
+
+    # Normalize config parameters from pixel space to normalized space
+    normalized_config = config.copy()
+    if 'min_distance' in config:
+        normalized_config['min_distance'] = config['min_distance'] / scene_size
+    if 'visual_range' in config:
+        normalized_config['visual_range'] = config['visual_range'] / scene_size
+
+    # Add 'independent' flag if not present
+    if 'independent' not in normalized_config:
+        normalized_config['independent'] = False
+
+    # Flatten for iteration
+    delta_vels = np.stack([VX.flatten(), VY.flatten()], axis=-1)  # [N, 2]
+    n_points = len(delta_vels)
+
+    # Compute forces at each velocity grid point
+    forces = np.zeros_like(delta_vels)
+
+    for idx in range(n_points):
+        # Both boids at the same position (square center: 0.5, 0.5 in normalized space)
+        # Boid i: stationary reference
+        boid_i = {
+            'x': 0.5,
+            'y': 0.5,
+            'dx': 0.0,
+            'dy': 0.0,
+            'species': 'A'
+        }
+
+        # Boid j: has relative velocity
+        vel_j = delta_vels[idx]
+        boid_j = {
+            'x': 0.5,  # Same position as i
+            'y': 0.5,
+            'dx': vel_j[0],
+            'dy': vel_j[1],
+            'species': 'A'
+        }
+
+        # Initial velocity of boid_j
+        initial_vel = np.array([boid_j['dx'], boid_j['dy']])
+
+        # Compute force ON boid_j due TO boid_i
+        # Put boids in a list and call force functions
+        boids = [boid_j, boid_i]
+
+        fly_towards_center(boid_j, boids, normalized_config)  # Rule 1: Cohesion
+        avoid_others(boid_j, boids, normalized_config)        # Rule 2: Separation
+        match_velocity(boid_j, boids, normalized_config)      # Rule 3: Alignment
+
+        # Force is change in velocity
+        final_vel = np.array([boid_j['dx'], boid_j['dy']])
+        forces[idx] = final_vel - initial_vel
+
+    # Reshape to grid
+    forces_x = forces[:, 0].reshape(grid_size, grid_size)
+    forces_y = forces[:, 1].reshape(grid_size, grid_size)
+    forces_mag = np.sqrt(forces_x**2 + forces_y**2)
+
+    return {
+        'VX': VX,
+        'VY': VY,
+        'forces': forces.reshape(grid_size, grid_size, 2),
+        'forces_x': forces_x,
+        'forces_y': forces_y,
+        'force_mag': forces_mag
+    }
+
+
+def compare_with_true_boids(
+    model,
+    save_path=None,
+    config=None,
+    particle_idx=0,
+    scene_size=480.0
+):
+    """
+    Compare learned interaction functions with true 2D boid rules.
+
+    Parameters
+    ----------
+    model : InteractionParticle
+        Trained model
+    save_path : str, optional
+        Path to save the plot
+    config : dict, optional
+        2D boid configuration dict with parameters:
+        - visual_range: 50.0 (pixels)
+        - min_distance: 15.0 (pixels)
+        - avoid_factor: 0.05
+        - matching_factor: 0.5
+        - centering_factor: 0.005
+    particle_idx : int
+        Particle index to use for embedding
+    scene_size : float
+        Scene size in pixels (default: 480 for 2D boids)
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        Figure object
+    """
+    # Default 2D boid parameters from boids_gnn_temp/boid.py
+    if config is None:
+        config = {
+            'visual_range': 50.0,         # Visual range for alignment and cohesion (pixels)
+            'min_distance': 15.0,         # Separation distance threshold (pixels)
+            'avoid_factor': 0.05,         # Separation weight
+            'matching_factor': 0.5,       # Alignment weight
+            'centering_factor': 0.005,    # Cohesion weight
+        }
+
+    # Create distance arrays
+    distances_fine = np.linspace(0.01, 1.0, 500)  # Normalized distances
+    distances_real = distances_fine * model.max_radius  # Real distances
+
+    # Get learned forces using evaluate_interaction
+    # Create delta_pos array along x-axis (particle j to the right of particle i)
+    with torch.no_grad():
+        delta_pos = torch.zeros((len(distances_fine), 2), dtype=torch.float32, device=model.device)
+        delta_pos[:, 0] = torch.tensor(distances_fine, dtype=torch.float32)  # Position along x-axis
+
+        # Use zero relative velocity (position-only forces)
+        delta_vel = torch.zeros_like(delta_pos)
+
+        # Evaluate interaction
+        learned_forces = model.evaluate_interaction(delta_pos, delta_vel, embedding_idx=particle_idx).cpu().numpy()
+
+    # Compute magnitude of learned forces
+    learned_magnitude = np.sqrt(learned_forces[:, 0]**2 + learned_forces[:, 1]**2)
+
+    # Convert distances to pixels
+    # 2D boids data is normalized by scene size (typically 480x480)
+    distances_pixels = distances_fine * scene_size
+
+    # Compute true 2D boid forces
+    true_separation = boid_separation_force(
+        distances_pixels,
+        min_distance=config['min_distance'],
+        avoid_factor=config['avoid_factor']
+    )
+    true_alignment = boid_alignment_force(
+        distances_pixels,
+        visual_range=config['visual_range'],
+        matching_factor=config['matching_factor']
+    )
+    true_cohesion = boid_cohesion_force(
+        distances_pixels,
+        visual_range=config['visual_range'],
+        centering_factor=config['centering_factor']
+    )
+
+    # Create comprehensive comparison plot
+    fig = plt.figure(figsize=(16, 10))
+    gs = fig.add_gridspec(3, 2, hspace=0.3, wspace=0.3)
+
+    # Plot 1: Learned force magnitude
+    ax1 = fig.add_subplot(gs[0, 0])
+    ax1.plot(distances_pixels, learned_magnitude, 'b-', linewidth=2, label='Learned Force Magnitude')
+    ax1.axhline(y=0, color='k', linestyle='--', alpha=0.3)
+    ax1.set_xlabel('Distance (pixels)', fontsize=11)
+    ax1.set_ylabel('Force Magnitude', fontsize=11)
+    ax1.set_title('Learned Interaction Force Magnitude', fontsize=12, fontweight='bold')
+    ax1.grid(True, alpha=0.3)
+    ax1.legend()
+
+    # Plot 2: Learned force components
+    ax2 = fig.add_subplot(gs[0, 1])
+    ax2.plot(distances_pixels, learned_forces[:, 0], 'b-', linewidth=2, label='X Component', alpha=0.7)
+    ax2.plot(distances_pixels, learned_forces[:, 1], 'r-', linewidth=2, label='Y Component', alpha=0.7)
+    ax2.axhline(y=0, color='k', linestyle='--', alpha=0.3)
+    ax2.set_xlabel('Distance (pixels)', fontsize=11)
+    ax2.set_ylabel('Force', fontsize=11)
+    ax2.set_title('Learned Force Components', fontsize=12, fontweight='bold')
+    ax2.grid(True, alpha=0.3)
+    ax2.legend()
+
+    # Plot 3: True boid separation (linear repulsion)
+    ax3 = fig.add_subplot(gs[1, 0])
+    ax3.plot(distances_pixels, true_separation, 'g-', linewidth=2, label='Separation Force')
+    ax3.axvline(x=config['min_distance'], color='g', linestyle='--', alpha=0.5,
+                label=f"min_distance = {config['min_distance']:.0f}px")
+    ax3.axhline(y=0, color='k', linestyle='--', alpha=0.3)
+    ax3.set_xlabel('Distance (pixels)', fontsize=11)
+    ax3.set_ylabel('Force', fontsize=11)
+    ax3.set_title('True Boid Separation (Linear Repulsion)', fontsize=12, fontweight='bold')
+    ax3.grid(True, alpha=0.3)
+    ax3.legend()
+    ax3.set_xlim([0, config['visual_range'] * 1.5])
+
+    # Plot 4: True boid alignment (step function)
+    ax4 = fig.add_subplot(gs[1, 1])
+    ax4.plot(distances_pixels, true_alignment, 'm-', linewidth=2, label='Alignment Weight')
+    ax4.axvline(x=config['visual_range'], color='m', linestyle='--', alpha=0.5,
+                label=f"visual_range = {config['visual_range']:.0f}px")
+    ax4.axhline(y=0, color='k', linestyle='--', alpha=0.3)
+    ax4.set_xlabel('Distance (pixels)', fontsize=11)
+    ax4.set_ylabel('Weight', fontsize=11)
+    ax4.set_title('True Boid Alignment (Step Function)', fontsize=12, fontweight='bold')
+    ax4.grid(True, alpha=0.3)
+    ax4.legend()
+    ax4.set_xlim([0, config['visual_range'] * 1.5])
+
+    # Plot 5: True boid cohesion (linear attraction)
+    ax5 = fig.add_subplot(gs[2, 0])
+    ax5.plot(distances_pixels, true_cohesion, 'c-', linewidth=2, label='Cohesion Force')
+    ax5.axvline(x=config['visual_range'], color='c', linestyle='--', alpha=0.5,
+                label=f"visual_range = {config['visual_range']:.0f}px")
+    ax5.axhline(y=0, color='k', linestyle='--', alpha=0.3)
+    ax5.set_xlabel('Distance (pixels)', fontsize=11)
+    ax5.set_ylabel('Force', fontsize=11)
+    ax5.set_title('True Boid Cohesion (Linear Attraction)', fontsize=12, fontweight='bold')
+    ax5.grid(True, alpha=0.3)
+    ax5.legend()
+    ax5.set_xlim([0, config['visual_range'] * 1.5])
+
+    # Plot 6: Combined comparison - overlay learned on true
+    ax6 = fig.add_subplot(gs[2, 1])
+
+    # Normalize learned force for comparison
+    learned_norm = learned_magnitude / (np.max(learned_magnitude) + 1e-8)
+
+    # Normalize true forces for comparison
+    sep_norm = true_separation / (np.max(true_separation) + 1e-8)
+    coh_norm = true_cohesion / (np.max(true_cohesion) + 1e-8)
+
+    ax6.plot(distances_pixels, learned_norm, 'b-', linewidth=2.5, label='Learned (normalized)', alpha=0.8)
+    ax6.plot(distances_pixels, sep_norm, 'g--', linewidth=2, label='Separation (normalized)', alpha=0.7)
+    ax6.plot(distances_pixels, coh_norm, 'c--', linewidth=2, label='Cohesion (normalized)', alpha=0.7)
+    ax6.axvline(x=config['min_distance'], color='g', linestyle=':', alpha=0.5)
+    ax6.axvline(x=config['visual_range'], color='c', linestyle=':', alpha=0.5)
+    ax6.set_xlabel('Distance (pixels)', fontsize=11)
+    ax6.set_ylabel('Normalized Force', fontsize=11)
+    ax6.set_title('Comparison: Learned vs True 2D Boid Rules', fontsize=12, fontweight='bold')
+    ax6.grid(True, alpha=0.3)
+    ax6.legend()
+    ax6.set_xlim([0, config['visual_range'] * 1.5])
+
+    # Add main title
+    fig.suptitle('InteractionParticle Model: Learned vs True 2D Boid Interaction Functions',
+                 fontsize=14, fontweight='bold', y=0.995)
+
+    # Add text box with parameters
+    param_text = (
+        f"2D Boid Parameters:\n"
+        f"visual_range = {config['visual_range']:.1f} px\n"
+        f"min_distance = {config['min_distance']:.1f} px\n"
+        f"avoid_factor = {config['avoid_factor']:.3f}\n"
+        f"matching_factor = {config['matching_factor']:.2f}\n"
+        f"centering_factor = {config['centering_factor']:.4f}"
+    )
+    fig.text(0.02, 0.98, param_text, transform=fig.transFigure,
+             fontsize=9, verticalalignment='top',
+             bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        logger.info(f"Saved comparison plot to {save_path}")
+
+    return fig
+
+
+def plot_symmetric_force_decomposition(position_forces, velocity_forces, save_path=None,
+                                       title_prefix="Learned", visual_range=None):
+    """
+    Plot symmetric force decomposition: position-dependent and velocity-dependent forces.
+
+    Creates a 1x2 grid:
+    - Left: Position-dependent forces (velocity=0, vary position)
+    - Right: Velocity-dependent forces (position=0, vary velocity)
+
+    Parameters
+    ----------
+    position_forces : dict
+        Output from evaluate_forces_on_grid (with delta_vel=0)
+    velocity_forces : dict
+        Output from evaluate_velocity_forces_on_grid (with delta_pos=0)
+    save_path : str, optional
+        Path to save figure
+    title_prefix : str
+        Prefix for title (e.g., "Learned" or "True Boid")
+    visual_range : float, optional
+        Visual range radius to draw as reference circle
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+    """
+    fig, axes = plt.subplots(1, 2, figsize=(16, 7))
+
+    # LEFT: Position-dependent forces (velocity = 0)
+    ax = axes[0]
+    X, Y = position_forces['X'], position_forces['Y']
+
+    # Plot force magnitude as heatmap
+    im = ax.pcolormesh(X, Y, position_forces['away_pos']['mag'],
+                       cmap='viridis', shading='auto')
+    plt.colorbar(im, ax=ax, label='Force Magnitude')
+
+    # Overlay force vectors (subsample for clarity)
+    step = max(1, len(X) // 15)
+    ax.quiver(X[::step, ::step], Y[::step, ::step],
+              position_forces['away_pos']['x'][::step, ::step],
+              position_forces['away_pos']['y'][::step, ::step],
+              color='white', alpha=0.7, scale=None, scale_units='xy')
+
+    ax.set_xlabel('Δx (relative position)')
+    ax.set_ylabel('Δy (relative position)')
+    ax.set_title(f'{title_prefix}: Position-Dependent Forces\n(Δv = 0)')
+    ax.set_aspect('equal')
+    ax.axhline(0, color='white', linestyle='--', alpha=0.3, linewidth=0.5)
+    ax.axvline(0, color='white', linestyle='--', alpha=0.3, linewidth=0.5)
+
+    # Add visual range circle if provided
+    if visual_range is not None:
+        circle = plt.Circle((0, 0), visual_range, fill=False,
+                           edgecolor='red', linestyle='--', linewidth=2,
+                           label=f'Visual range = {visual_range:.2f}')
+        ax.add_patch(circle)
+        ax.legend(loc='upper right')
+
+    # RIGHT: Velocity-dependent forces (position = 0)
+    ax = axes[1]
+    VX, VY = velocity_forces['VX'], velocity_forces['VY']
+
+    # Plot force magnitude as heatmap
+    im = ax.pcolormesh(VX, VY, velocity_forces['force_mag'],
+                       cmap='plasma', shading='auto')
+    plt.colorbar(im, ax=ax, label='Force Magnitude')
+
+    # Overlay force vectors (subsample for clarity)
+    step = max(1, len(VX) // 15)
+    ax.quiver(VX[::step, ::step], VY[::step, ::step],
+              velocity_forces['forces_x'][::step, ::step],
+              velocity_forces['forces_y'][::step, ::step],
+              color='white', alpha=0.7, scale=None, scale_units='xy')
+
+    ax.set_xlabel('Δv_x (relative velocity)')
+    ax.set_ylabel('Δv_y (relative velocity)')
+    ax.set_title(f'{title_prefix}: Velocity-Dependent Forces\n(Δpos = 0)')
+    ax.set_aspect('equal')
+    ax.axhline(0, color='white', linestyle='--', alpha=0.3, linewidth=0.5)
+    ax.axvline(0, color='white', linestyle='--', alpha=0.3, linewidth=0.5)
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        logger.info(f"Saved symmetric force decomposition plot to {save_path}")
+
+    return fig
+
+
+def plot_training_history(history, save_path=None):
+    """
+    Plot training and validation loss curves.
+
+    Parameters
+    ----------
+    history : dict
+        Training history with 'train_loss' and 'val_loss' keys
+    save_path : str, optional
+        Path to save the plot
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        Figure object
+    """
+    fig, ax = plt.subplots(1, 1, figsize=(10, 6))
+
+    epochs = range(1, len(history['train_loss']) + 1)
+    ax.plot(epochs, history['train_loss'], 'b-', linewidth=2, label='Train Loss')
+    ax.plot(epochs, history['val_loss'], 'r-', linewidth=2, label='Val Loss')
+    ax.set_xlabel('Epoch', fontsize=12)
+    ax.set_ylabel('Loss (MSE)', fontsize=12)
+    ax.set_title('Training History', fontsize=14, fontweight='bold')
+    ax.grid(True, alpha=0.3)
+    ax.legend()
+    ax.set_yscale('log')
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        logger.info(f"Saved training history plot to {save_path}")
+
+    return fig
+
+
+def plot_rollout_comparison(ground_truth, predicted, trajectory_idx=0, n_particles=None, save_path=None):
+    """
+    Create side-by-side comparison of ground truth vs predicted trajectories.
+
+    Parameters
+    ----------
+    ground_truth : np.ndarray
+        Ground truth positions [T, N, 2]
+    predicted : np.ndarray
+        Predicted positions [T, N, 2]
+    trajectory_idx : int
+        Index identifier for the trajectory
+    n_particles : int, optional
+        Number of particles to plot (default: all)
+    save_path : str, optional
+        Path to save figure
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        Figure object
+    """
+    T, N, D = ground_truth.shape
+
+    if n_particles is None:
+        n_particles = N
+    else:
+        n_particles = min(n_particles, N)
+
+    # Create figure with two subplots side by side
+    fig, axes = plt.subplots(1, 2, figsize=(16, 7))
+
+    # Plot ground truth
+    ax = axes[0]
+    for i in range(n_particles):
+        ax.plot(ground_truth[:, i, 0], ground_truth[:, i, 1],
+                alpha=0.6, linewidth=1.5)
+        # Mark start and end
+        ax.scatter(ground_truth[0, i, 0], ground_truth[0, i, 1],
+                  c='green', s=100, marker='o', zorder=5, edgecolors='black', linewidth=1)
+        ax.scatter(ground_truth[-1, i, 0], ground_truth[-1, i, 1],
+                  c='red', s=100, marker='s', zorder=5, edgecolors='black', linewidth=1)
+
+        # Add initial velocity vector
+        if T >= 2:
+            initial_vel = ground_truth[1, i] - ground_truth[0, i]
+            # Scale for visibility (multiply by 3 for better visualization)
+            scale = 3.0
+            ax.arrow(ground_truth[0, i, 0], ground_truth[0, i, 1],
+                    initial_vel[0] * scale, initial_vel[1] * scale,
+                    head_width=0.02, head_length=0.015, fc='blue', ec='blue',
+                    alpha=0.7, linewidth=1.5, zorder=4)
+
+    ax.set_xlabel('X Position', fontsize=12)
+    ax.set_ylabel('Y Position', fontsize=12)
+    ax.set_title('Ground Truth', fontsize=14, fontweight='bold')
+    ax.grid(True, alpha=0.3)
+    ax.set_aspect('equal')
+    ax.set_xlim([-0.3, 1.3])
+    ax.set_ylim([-0.3, 1.3])
+
+    # Plot predicted
+    ax = axes[1]
+    for i in range(n_particles):
+        ax.plot(predicted[:, i, 0], predicted[:, i, 1],
+                alpha=0.6, linewidth=1.5)
+        # Mark start and end
+        ax.scatter(predicted[0, i, 0], predicted[0, i, 1],
+                  c='green', s=100, marker='o', zorder=5, edgecolors='black', linewidth=1)
+        ax.scatter(predicted[-1, i, 0], predicted[-1, i, 1],
+                  c='red', s=100, marker='s', zorder=5, edgecolors='black', linewidth=1)
+
+        # Add initial velocity vector (same initial velocity as ground truth)
+        if T >= 2:
+            initial_vel = ground_truth[1, i] - ground_truth[0, i]
+            # Scale for visibility (multiply by 3 for better visualization)
+            scale = 3.0
+            ax.arrow(predicted[0, i, 0], predicted[0, i, 1],
+                    initial_vel[0] * scale, initial_vel[1] * scale,
+                    head_width=0.02, head_length=0.015, fc='blue', ec='blue',
+                    alpha=0.7, linewidth=1.5, zorder=4)
+
+    ax.set_xlabel('X Position', fontsize=12)
+    ax.set_ylabel('Y Position', fontsize=12)
+    ax.set_title('Predicted (Model)', fontsize=14, fontweight='bold')
+    ax.grid(True, alpha=0.3)
+    ax.set_aspect('equal')
+    ax.set_xlim([-0.3, 1.3])
+    ax.set_ylim([-0.3, 1.3])
+
+    # Add legend
+    from matplotlib.patches import Patch, FancyArrow
+    legend_elements = [
+        Patch(facecolor='green', edgecolor='black', label='Start'),
+        Patch(facecolor='red', edgecolor='black', label='End'),
+        Patch(facecolor='blue', edgecolor='blue', label='Initial Velocity (3x scaled)')
+    ]
+    fig.legend(handles=legend_elements, loc='upper center', ncol=3, fontsize=11)
+
+    plt.suptitle(f'Trajectory Comparison (Sample {trajectory_idx})',
+                 fontsize=16, fontweight='bold', y=0.98)
+    plt.tight_layout(rect=[0, 0, 1, 0.96])
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        logger.info(f"Saved rollout comparison to {save_path}")
+
+    return fig
+
+
+def plot_rollout_error_over_time(ground_truth_list, predicted_list, save_path=None):
+    """
+    Plot position error as a function of time for multiple rollouts.
+
+    Parameters
+    ----------
+    ground_truth_list : list of np.ndarray
+        List of ground truth trajectories, each [T, N, 2]
+    predicted_list : list of np.ndarray
+        List of predicted trajectories, each [T, N, 2]
+    save_path : str, optional
+        Path to save figure
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        Figure object
+    """
+    # Compute errors for each trajectory
+    errors_over_time = []
+
+    for gt, pred in zip(ground_truth_list, predicted_list):
+        # Compute per-particle error at each timestep
+        error = np.linalg.norm(gt - pred, axis=-1)  # [T, N]
+        # Average over particles
+        mean_error = np.mean(error, axis=1)  # [T]
+        errors_over_time.append(mean_error)
+
+    # Convert to array
+    errors_over_time = np.array(errors_over_time)  # [n_trajectories, T]
+
+    # Compute statistics
+    mean_error = np.mean(errors_over_time, axis=0)
+    std_error = np.std(errors_over_time, axis=0)
+
+    # Create figure
+    fig, ax = plt.subplots(1, 1, figsize=(10, 6))
+
+    timesteps = np.arange(len(mean_error))
+
+    # Plot mean and std
+    ax.plot(timesteps, mean_error, 'b-', linewidth=2, label='Mean Error')
+    ax.fill_between(timesteps,
+                     mean_error - std_error,
+                     mean_error + std_error,
+                     alpha=0.3, color='b', label='±1 Std')
+
+    ax.set_xlabel('Time Step', fontsize=12)
+    ax.set_ylabel('Position Error', fontsize=12)
+    ax.set_title('Rollout Position Error Over Time', fontsize=14, fontweight='bold')
+    ax.grid(True, alpha=0.3)
+    ax.legend()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        logger.info(f"Saved error over time plot to {save_path}")
+
+    return fig
+
+
+def create_rollout_report(results, save_dir=None):
+    """
+    Create comprehensive rollout evaluation report with all visualizations.
+
+    Parameters
+    ----------
+    results : dict
+        Results from evaluate_rollout()
+    save_dir : str, optional
+        Directory to save figures
+
+    Returns
+    -------
+    figures : dict
+        Dictionary of matplotlib figures
+    """
+    import os
+
+    figures = {}
+
+    # Extract data
+    gt_positions = results['ground_truth_positions']
+    pred_positions = results['predicted_positions']
+    metrics = results['metrics']
+
+    logger.info(f"Creating rollout report for {len(gt_positions)} trajectories")
+
+    # 1. Plot first few trajectory comparisons
+    n_examples = min(3, len(gt_positions))
+    for i in range(n_examples):
+        fig = plot_rollout_comparison(
+            gt_positions[i],
+            pred_positions[i],
+            trajectory_idx=i,
+            save_path=os.path.join(save_dir, f'rollout_comparison_{i}.png') if save_dir else None
+        )
+        figures[f'comparison_{i}'] = fig
+        plt.close(fig)
+
+    # 2. Plot error over time
+    fig = plot_rollout_error_over_time(
+        gt_positions,
+        pred_positions,
+        save_path=os.path.join(save_dir, 'rollout_error_over_time.png') if save_dir else None
+    )
+    figures['error_over_time'] = fig
+    plt.close(fig)
+
+    # 3. Print metrics summary
+    logger.info("=" * 60)
+    logger.info("ROLLOUT EVALUATION METRICS")
+    logger.info("=" * 60)
+    logger.info(f"Number of trajectories: {metrics['n_trajectories']}")
+    logger.info(f"Mean position error: {metrics['mean_position_error']:.6f} ± {metrics['std_position_error']:.6f}")
+    logger.info(f"Mean velocity error: {metrics['mean_velocity_error']:.6f} ± {metrics['std_velocity_error']:.6f}")
+    logger.info("=" * 60)
+
+    return figures
diff --git a/collab_env/gnn/interaction_particles/run_training.py b/collab_env/gnn/interaction_particles/run_training.py
new file mode 100644
index 0000000..703ab5a
--- /dev/null
+++ b/collab_env/gnn/interaction_particles/run_training.py
@@ -0,0 +1,696 @@
+#!/usr/bin/env python
+"""
+Main script for training InteractionParticle model on boids data.
+
+Usage:
+    python -m collab_env.gnn.interaction_particles.run_training [options]
+
+Example:
+    python -m collab_env.gnn.interaction_particles.run_training \
+        --dataset collab_env/data/boids/boid_single_species_basic.pt \
+        --epochs 100 \
+        --batch-size 32 \
+        --save-dir trained_models/interaction_particle
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+import torch
+import yaml
+from loguru import logger
+
+from .train import train_interaction_particle, evaluate_model, evaluate_rollout, compute_velocity_statistics
+from .plotting import (
+    plot_training_history,
+    create_rollout_report,
+    evaluate_forces_on_grid,
+    evaluate_velocity_forces_on_grid,
+    plot_force_decomposition,
+    plot_symmetric_force_decomposition,
+    evaluate_true_boid_forces,
+    evaluate_true_boid_velocity_forces,
+    plot_rollout_comparison
+)
+from collab_env.data.file_utils import expand_path
+
+
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description='Train InteractionParticle model on boids data',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    # Data arguments
+    parser.add_argument(
+        '--dataset',
+        type=str,
+        default='simulated_data/boid_single_species_basic.pt',
+        help='Path to 2D boids dataset .pt file (AnimalTrajectoryDataset)'
+    )
+    parser.add_argument(
+        '--config',
+        type=str,
+        default=None,
+        help='Path to 2D boids config file (.pt). If not provided, auto-detects or uses defaults.'
+    )
+
+    # Training arguments
+    parser.add_argument('--epochs', type=int, default=100, help='Number of training epochs')
+    parser.add_argument('--batch-size', type=int, default=32, help='Batch size')
+    parser.add_argument('--learning-rate', type=float, default=1e-5, help='Learning rate')
+    parser.add_argument('--train-split', type=float, default=0.8, help='Train/val split ratio')
+    parser.add_argument('--visual-range', type=float, default=0.3, help='Visual range for edge construction (normalized)')
+    parser.add_argument('--scene-size', type=float, default=480.0, help='Scene size in pixels (for normalizing config parameters)')
+    parser.add_argument('--rollout-steps', type=int, default=1, help='Number of rollout steps for training (1=single-step, >1=multi-step)')
+    parser.add_argument('--seed', type=int, default=42, help='Random seed')
+
+    # Model arguments
+    parser.add_argument('--hidden-dim', type=int, default=128, help='Hidden dimension for MLP')
+    parser.add_argument('--embedding-dim', type=int, default=16, help='Particle embedding dimension')
+    parser.add_argument('--n-layers', type=int, default=3, help='Number of MLP layers')
+    parser.add_argument('--n-particles', type=int, default=20, help='Number of particles (auto-detected from config if available)')
+
+    # Output arguments
+    parser.add_argument(
+        '--save-dir',
+        type=str,
+        default='trained_models/interaction_particle',
+        help='Directory to save model and plots'
+    )
+    parser.add_argument('--no-plot', action='store_true', help='Skip plotting')
+    parser.add_argument('--device', type=str, default=None, help='Device to use (cpu/cuda/cuda:0/mps/auto)')
+    parser.add_argument('--plot-every', type=int, default=0, help='Plot force decomposition and rollout every N epochs (0=only at end)')
+
+    # Evaluation
+    parser.add_argument('--eval-only', action='store_true', help='Only evaluate a saved model')
+    parser.add_argument('--model-path', type=str, default=None, help='Path to saved model for evaluation')
+    parser.add_argument('--evaluate-rollout', action='store_true', help='Evaluate model with multi-step rollout')
+    parser.add_argument('--n-rollout-steps', type=int, default=50, help='Number of steps for rollout evaluation')
+
+    return parser.parse_args()
+
+
+def load_boids_config(config_path, dataset_path : Path=None):
+    """
+    Load 2D boids config from .pt file.
+
+    The config file contains species_configs dictionary with parameters like:
+    - visual_range: 50.0
+    - min_distance: 15.0
+    - avoid_factor: 0.05
+    - matching_factor: 0.5
+    - centering_factor: 0.005
+    """
+    # Default 2D boid parameters
+    default_config = {
+        'visual_range': 50.0,
+        'min_distance': 15.0,
+        'avoid_factor': 0.05,
+        'matching_factor': 0.5,
+        'centering_factor': 0.005,
+    }
+
+    if config_path is None:
+        # Try to infer config from dataset path
+        if dataset_path:
+            config_path_guess = dataset_path.with_stem(dataset_path.stem + '_config')
+            if os.path.exists(expand_path(config_path_guess)):
+                config_path = config_path_guess
+                logger.info(f"Auto-detected config: {config_path}")
+            else:
+                logger.info("No config file found, using defaults")
+                return default_config
+        else:
+            logger.info("Using default 2D boid parameters")
+            return default_config
+
+    try:
+        config_path = expand_path(config_path)
+        # Load species_configs from .pt file
+        species_configs = torch.load(config_path, weights_only=False)
+        # Extract first species config (usually 'A')
+        first_species = list(species_configs.keys())[0]
+        config = species_configs[first_species]
+        logger.info(f"Loaded config for species '{first_species}'")
+        return config
+
+    except Exception as e:
+        logger.warning(f"Could not load config from {config_path}: {e}")
+        logger.warning("Using default parameters")
+        return default_config
+
+
+def main():
+    """Main function."""
+    args = parse_args()
+
+    # Setup logging
+    logger.remove()
+    logger.add(sys.stderr, level="INFO")
+
+    # Expand paths
+    dataset_path = expand_path(args.dataset)
+    save_dir = expand_path(args.save_dir)
+    os.makedirs(save_dir, exist_ok=True)
+
+    # Add file logging
+    log_path = os.path.join(save_dir, 'training.log')
+    logger.add(log_path, level="DEBUG")
+
+    logger.info("=" * 60)
+    logger.info("InteractionParticle Training on 2D Boids Data")
+    logger.info("=" * 60)
+    logger.info(f"Dataset: {dataset_path}")
+    logger.info(f"Save directory: {save_dir}")
+    logger.info(f"Device: {args.device or 'auto'}")
+
+    # Load boids config
+    boids_config = load_boids_config(args.config, dataset_path)
+    if boids_config:
+        logger.info("Boids config parameters:")
+        for key, value in boids_config.items():
+            logger.info(f"  {key}: {value}")
+    else:
+        logger.info("No boids config loaded")
+
+    # Compute velocity statistics from dataset for visualization
+    logger.info("Computing velocity statistics from dataset...")
+    vel_stats = compute_velocity_statistics(dataset_path)
+
+    # Use for speed magnitude in position-dependent force plots
+    speed_magnitude = vel_stats['mean_speed']
+    logger.info(f"Using mean speed = {speed_magnitude:.4f} for position-dependent force visualization")
+
+    # Use for velocity grid bounds in velocity-dependent force plots
+    max_rel_vel = vel_stats['max_relative_velocity']
+    logger.info(f"Using max relative velocity = {max_rel_vel:.4f} for velocity-dependent force visualization")
+
+    # Set device
+    # Device selection with MPS support
+    if args.device and args.device != 'auto':
+        device = torch.device(args.device)
+    else:
+        # Auto-detect: cuda > mps > cpu
+        if torch.cuda.is_available():
+            device = torch.device('cuda')
+        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+            device = torch.device('mps')
+        else:
+            device = torch.device('cpu')
+    logger.info(f"Using device: {device}")
+
+    # Evaluation mode
+    if args.eval_only:
+        if args.model_path is None:
+            logger.error("Must specify --model-path for evaluation mode")
+            sys.exit(1)
+
+        logger.info("Loading model for evaluation...")
+        checkpoint = torch.load(args.model_path, map_location=device)
+
+        # Create model
+        from .model import InteractionParticle
+        model = InteractionParticle(checkpoint['config'], device=device)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        model.eval()
+
+        logger.info("Evaluating model...")
+        metrics = evaluate_model(model, dataset_path, visual_range=args.visual_range, device=device, train_split=args.train_split)
+        logger.info(f"Evaluation metrics: {metrics}")
+
+        # Generate plots (same as post-training)
+        if not args.no_plot:
+            logger.info("\nGenerating plots...")
+
+            # Generate 2D force decomposition plots (2x3 layout)
+            logger.info("Generating 2D force decomposition plots...")
+
+            # Use the training visual_range parameter for BOTH learned and true forces
+            # This ensures fair comparison (same plot range for both)
+            plot_max_dist = args.visual_range * 1.5  # 50% cushion
+
+            # Evaluate learned forces on grid (use same speed magnitude as true forces)
+            learned_forces = evaluate_forces_on_grid(
+                model,
+                grid_size=50,
+                max_dist=plot_max_dist,
+                particle_idx=0,
+                speed_magnitude=speed_magnitude
+            )
+
+            # Evaluate ground truth boid forces (use SAME range as learned for comparison)
+            true_forces = evaluate_true_boid_forces(
+                boids_config,
+                grid_size=50,
+                max_dist=plot_max_dist,
+                scene_size=args.scene_size
+            )
+
+            # Plot learned forces
+            learned_plot_path = os.path.join(save_dir, 'learned_force_decomposition.png')
+            plot_force_decomposition(
+                learned_forces,
+                save_path=learned_plot_path,
+                title_prefix="Learned",
+                visual_range=args.visual_range
+            )
+
+            # Plot ground truth forces
+            true_plot_path = os.path.join(save_dir, 'true_boid_force_decomposition.png')
+            # Extract min_distance from boids_config (normalized to [0,1] space)
+            min_distance_normalized = None
+            if boids_config and 'min_distance' in boids_config:
+                min_distance_normalized = boids_config['min_distance'] / args.scene_size
+
+            plot_force_decomposition(
+                true_forces,
+                save_path=true_plot_path,
+                title_prefix="True Boid",
+                visual_range=args.visual_range,
+                min_distance=min_distance_normalized
+            )
+
+            # Generate SYMMETRIC force decomposition plots
+            logger.info("Generating symmetric force decomposition plots...")
+
+            # Position-dependent forces (with delta_vel = 0)
+            pos_forces_learned = evaluate_forces_on_grid(
+                model,
+                grid_size=50,
+                max_dist=plot_max_dist,
+                particle_idx=0,
+                speed_magnitude=0.0  # Zero velocity for position-only term
+            )
+
+            # Velocity-dependent forces (with delta_pos = 0)
+            vel_forces_learned = evaluate_velocity_forces_on_grid(
+                model,
+                grid_size=50,
+                max_vel=max_rel_vel,
+                particle_idx=0
+            )
+
+            # Plot symmetric decomposition for learned forces
+            symmetric_learned_path = os.path.join(save_dir, 'learned_symmetric_decomposition.png')
+            plot_symmetric_force_decomposition(
+                pos_forces_learned,
+                vel_forces_learned,
+                save_path=symmetric_learned_path,
+                title_prefix="Learned",
+                visual_range=args.visual_range
+            )
+
+            # Generate TRUE BOID symmetric decomposition for comparison
+            logger.info("Generating true boid symmetric force decomposition...")
+
+            # Position-dependent forces (true boids, delta_vel = 0)
+            # Reuse the position grid from true_forces (already computed with vel=0)
+            pos_forces_true = true_forces  # This has away_pos which is computed with vel=0
+
+            # Velocity-dependent forces (true boids, delta_pos = 0)
+            vel_forces_true = evaluate_true_boid_velocity_forces(
+                boids_config,
+                grid_size=50,
+                max_vel=max_rel_vel,
+                scene_size=args.scene_size
+            )
+
+            # Plot symmetric decomposition for true boid forces
+            symmetric_true_path = os.path.join(save_dir, 'true_boid_symmetric_decomposition.png')
+            plot_symmetric_force_decomposition(
+                pos_forces_true,
+                vel_forces_true,
+                save_path=symmetric_true_path,
+                title_prefix="True Boid",
+                visual_range=args.visual_range
+            )
+
+            logger.info(f"Plots saved to {save_dir}")
+
+        # Rollout evaluation (if requested)
+        if args.evaluate_rollout:
+            logger.info("\nEvaluating model with multi-step rollout...")
+            rollout_results = evaluate_rollout(
+                model,
+                dataset_path,
+                visual_range=args.visual_range,
+                n_rollout_steps=args.n_rollout_steps,
+                device=device,
+                train_split=args.train_split
+            )
+
+            # Create rollout report
+            rollout_dir = os.path.join(save_dir, 'rollout_evaluation')
+            os.makedirs(rollout_dir, exist_ok=True)
+
+            logger.info(f"Creating rollout visualizations...")
+            create_rollout_report(rollout_results, save_dir=rollout_dir)
+
+            logger.info(f"Rollout evaluation saved to {rollout_dir}")
+
+        return
+
+    # Determine n_particles from boids config or command line
+    if boids_config and 'counts' in boids_config:
+        n_particles = boids_config['counts']
+        logger.info(f"Using n_particles={n_particles} from boids config")
+    else:
+        n_particles = args.n_particles
+        logger.info(f"Using n_particles={n_particles} from command line argument")
+
+    # Create model config
+    model_config = {
+        'n_particles': n_particles,
+        'n_particle_types': 1,
+        'max_radius': 1.0,
+        'hidden_dim': args.hidden_dim,
+        'embedding_dim': args.embedding_dim,
+        'n_mp_layers': args.n_layers,
+        'input_size': 11,  # delta_pos(2) + r(1) + delta_vel(2) + pos_i(2) + vel_i(2) + boundary_complement(2)
+        'output_size': 2,
+    }
+
+    logger.info("Model configuration:")
+    for key, value in model_config.items():
+        logger.info(f"  {key}: {value}")
+
+    # Define epoch callback for plotting during training
+    def epoch_plotting_callback(model, epoch, train_loss, val_loss, save_dir):
+        """Generate plots after each validation epoch if plot_every is set."""
+        if args.plot_every <= 0 or epoch % args.plot_every != 0:
+            return
+
+        logger.info(f"\nGenerating epoch {epoch} plots...")
+        epoch_dir = os.path.join(save_dir, f'epoch_{epoch:04d}')
+        os.makedirs(epoch_dir, exist_ok=True)
+
+        # Plot force decomposition comparison
+        try:
+            # Use the training visual_range parameter for BOTH learned and true forces
+            # This ensures fair comparison (same plot range for both)
+            plot_max_dist = args.visual_range * 1.5  # 50% cushion
+
+            # Evaluate learned forces (use same speed magnitude as true forces)
+            learned_forces = evaluate_forces_on_grid(
+                model,
+                grid_size=50,
+                max_dist=plot_max_dist,
+                particle_idx=0,
+                speed_magnitude=speed_magnitude
+            )
+
+            # Evaluate ground truth forces (use SAME range as learned for comparison)
+            true_forces = evaluate_true_boid_forces(
+                boids_config,
+                grid_size=50,
+                max_dist=plot_max_dist,
+                scene_size=args.scene_size
+            )
+
+            # Plot learned forces
+            learned_plot_path = os.path.join(epoch_dir, 'learned_force_decomposition.png')
+            plot_force_decomposition(
+                learned_forces,
+                save_path=learned_plot_path,
+                title_prefix=f"Learned (Epoch {epoch})",
+                visual_range=args.visual_range
+            )
+
+            # Plot ground truth forces
+            true_plot_path = os.path.join(epoch_dir, 'true_boid_force_decomposition.png')
+            # Extract min_distance from boids_config (normalized to [0,1] space)
+            min_distance_normalized = None
+            if boids_config and 'min_distance' in boids_config:
+                min_distance_normalized = boids_config['min_distance'] / args.scene_size
+
+            plot_force_decomposition(
+                true_forces,
+                save_path=true_plot_path,
+                title_prefix=f"True Boid (Epoch {epoch})",
+                visual_range=args.visual_range,
+                min_distance=min_distance_normalized
+            )
+
+            logger.info(f"Saved force decomposition plots to {epoch_dir}")
+        except Exception as e:
+            logger.warning(f"Failed to generate force decomposition plots: {e}")
+
+        # Plot symmetric decomposition
+        try:
+            logger.info(f"Generating symmetric force decomposition for epoch {epoch}...")
+
+            # Compute velocity statistics for max_rel_vel
+            max_rel_vel = 0.05  # Use same default as after-training plots
+
+            # Evaluate learned forces on velocity grid (delta_pos = 0)
+            learned_vel_forces = evaluate_velocity_forces_on_grid(
+                model,
+                grid_size=50,
+                max_vel=max_rel_vel,
+                particle_idx=0
+            )
+
+            # Evaluate true boid forces on velocity grid (delta_pos = 0)
+            true_vel_forces = evaluate_true_boid_velocity_forces(
+                boids_config,
+                grid_size=50,
+                max_vel=max_rel_vel,
+                scene_size=args.scene_size
+            )
+
+            # Plot learned symmetric decomposition
+            learned_symmetric_path = os.path.join(epoch_dir, 'learned_symmetric_decomposition.png')
+            plot_symmetric_force_decomposition(
+                learned_forces,  # position-dependent (already computed above)
+                learned_vel_forces,  # velocity-dependent
+                save_path=learned_symmetric_path,
+                title_prefix=f"Learned (Epoch {epoch})",
+                visual_range=args.visual_range
+            )
+
+            # Plot true boid symmetric decomposition
+            true_symmetric_path = os.path.join(epoch_dir, 'true_boid_symmetric_decomposition.png')
+            plot_symmetric_force_decomposition(
+                true_forces,  # position-dependent (already computed above)
+                true_vel_forces,  # velocity-dependent
+                save_path=true_symmetric_path,
+                title_prefix=f"True Boid (Epoch {epoch})",
+                visual_range=args.visual_range
+            )
+
+            logger.info(f"Saved symmetric decomposition plots to {epoch_dir}")
+        except Exception as e:
+            logger.warning(f"Failed to generate symmetric decomposition plots: {e}")
+
+        # Plot rollout comparison
+        if args.evaluate_rollout:
+            try:
+                rollout_results = evaluate_rollout(
+                    model,
+                    dataset_path,
+                    visual_range=args.visual_range,
+                    n_rollout_steps=args.n_rollout_steps,
+                    device=device,
+                    train_split=args.train_split
+                )
+
+                # Save rollout comparison (use first trajectory)
+                if len(rollout_results['ground_truth_positions']) > 0:
+                    rollout_plot_path = os.path.join(epoch_dir, 'rollout_comparison.png')
+                    plot_rollout_comparison(
+                        rollout_results['ground_truth_positions'][0],
+                        rollout_results['predicted_positions'][0],
+                        trajectory_idx=0,
+                        save_path=rollout_plot_path
+                    )
+                else:
+                    logger.warning(f"No rollout trajectories to plot")
+
+                logger.info(f"Saved rollout comparison to {epoch_dir}")
+            except Exception as e:
+                logger.warning(f"Failed to generate rollout comparison: {e}")
+
+    # Train model
+    logger.info("\nStarting training...")
+    model, history = train_interaction_particle(
+        dataset_path=dataset_path,
+        config=model_config,
+        epochs=args.epochs,
+        batch_size=args.batch_size,
+        learning_rate=args.learning_rate,
+        train_split=args.train_split,
+        visual_range=args.visual_range,
+        rollout_steps=args.rollout_steps,
+        device=device,
+        save_dir=save_dir,
+        seed=args.seed,
+        epoch_callback=epoch_plotting_callback if args.plot_every > 0 else None
+    )
+
+    logger.info("\nTraining completed!")
+    logger.info(f"Final train loss: {history['train_loss'][-1]:.2g}")
+    logger.info(f"Final val loss: {history['val_loss'][-1]:.2g}")
+
+    # Save model info
+    info_path = os.path.join(save_dir, 'model_info.yaml')
+    with open(info_path, 'w') as f:
+        yaml.dump({
+            'args': vars(args),
+            'config': model_config,
+            'boids_config': boids_config,
+            'final_train_loss': float(history['train_loss'][-1]),
+            'final_val_loss': float(history['val_loss'][-1]),
+        }, f, default_flow_style=False)
+    logger.info(f"Saved model info to {info_path}")
+
+    # Generate plots
+    if not args.no_plot:
+        logger.info("\nGenerating plots...")
+
+        # Plot training history
+        history_path = os.path.join(save_dir, 'training_history.png')
+        plot_training_history(history, save_path=history_path)
+
+        # Generate 2D force decomposition plots (2x3 layout)
+        logger.info("Generating 2D force decomposition plots...")
+
+        # Use the training visual_range parameter for BOTH learned and true forces
+        # This ensures fair comparison (same plot range for both)
+        plot_max_dist = args.visual_range * 1.5  # 50% cushion
+
+        # Evaluate learned forces on grid (use same speed magnitude as true forces)
+        learned_forces = evaluate_forces_on_grid(
+            model,
+            grid_size=50,
+            max_dist=plot_max_dist,
+            particle_idx=0,
+            speed_magnitude=speed_magnitude
+        )
+
+        # Evaluate ground truth boid forces (use SAME range as learned for comparison)
+        true_forces = evaluate_true_boid_forces(
+            boids_config,
+            grid_size=50,
+            max_dist=plot_max_dist,
+            scene_size=args.scene_size
+        )
+
+        # Plot learned forces
+        learned_plot_path = os.path.join(save_dir, 'learned_force_decomposition.png')
+        plot_force_decomposition(
+            learned_forces,
+            save_path=learned_plot_path,
+            title_prefix="Learned",
+            visual_range=args.visual_range
+        )
+
+        # Plot ground truth forces
+        true_plot_path = os.path.join(save_dir, 'true_boid_force_decomposition.png')
+        # Extract min_distance from boids_config (normalized to [0,1] space)
+        min_distance_normalized = None
+        if boids_config and 'min_distance' in boids_config:
+            min_distance_normalized = boids_config['min_distance'] / args.scene_size
+
+        plot_force_decomposition(
+            true_forces,
+            save_path=true_plot_path,
+            title_prefix="True Boid",
+            visual_range=args.visual_range,
+            min_distance=min_distance_normalized
+        )
+
+        # Generate SYMMETRIC force decomposition plots (NEW!)
+        logger.info("Generating symmetric force decomposition plots...")
+
+        # Position-dependent forces (with delta_vel = 0)
+        pos_forces_learned = evaluate_forces_on_grid(
+            model,
+            grid_size=50,
+            max_dist=plot_max_dist,
+            particle_idx=0,
+            speed_magnitude=0.0  # Zero velocity for position-only term
+        )
+
+        # Velocity-dependent forces (with delta_pos = 0)
+        vel_forces_learned = evaluate_velocity_forces_on_grid(
+            model,
+            grid_size=50,
+            max_vel=max_rel_vel,
+            particle_idx=0
+        )
+
+        # Plot symmetric decomposition for learned forces
+        symmetric_learned_path = os.path.join(save_dir, 'learned_symmetric_decomposition.png')
+        plot_symmetric_force_decomposition(
+            pos_forces_learned,
+            vel_forces_learned,
+            save_path=symmetric_learned_path,
+            title_prefix="Learned",
+            visual_range=args.visual_range
+        )
+
+        # Generate TRUE BOID symmetric decomposition for comparison
+        logger.info("Generating true boid symmetric force decomposition...")
+
+        # Position-dependent forces (true boids, delta_vel = 0)
+        # Reuse the position grid from true_forces (already computed with vel=0)
+        pos_forces_true = true_forces  # This has away_pos which is computed with vel=0
+
+        # Velocity-dependent forces (true boids, delta_pos = 0)
+        vel_forces_true = evaluate_true_boid_velocity_forces(
+            boids_config,
+            grid_size=50,
+            max_vel=max_rel_vel,
+            scene_size=args.scene_size
+        )
+
+        # Plot symmetric decomposition for true boid forces
+        symmetric_true_path = os.path.join(save_dir, 'true_boid_symmetric_decomposition.png')
+        plot_symmetric_force_decomposition(
+            pos_forces_true,
+            vel_forces_true,
+            save_path=symmetric_true_path,
+            title_prefix="True Boid",
+            visual_range=args.visual_range
+        )
+
+        logger.info(f"Plots saved to {save_dir}")
+
+    # Rollout evaluation
+    if args.evaluate_rollout:
+        logger.info("\nEvaluating model with multi-step rollout...")
+        rollout_results = evaluate_rollout(
+            model,
+            dataset_path,
+            visual_range=args.visual_range,
+            n_rollout_steps=args.n_rollout_steps,
+            device=device,
+            train_split=args.train_split
+        )
+
+        # Create rollout report
+        rollout_dir = os.path.join(save_dir, 'rollout_evaluation')
+        os.makedirs(rollout_dir, exist_ok=True)
+
+        logger.info(f"Creating rollout visualizations...")
+        create_rollout_report(rollout_results, save_dir=rollout_dir)
+
+        logger.info(f"Rollout evaluation saved to {rollout_dir}")
+
+    logger.info("\n" + "=" * 60)
+    logger.info("Training pipeline completed successfully!")
+    logger.info("=" * 60)
+    logger.info(f"\nResults saved to: {save_dir}")
+    logger.info(f"  - Model: best_model.pt, final_model.pt")
+    logger.info(f"  - Plots: *.png")
+    logger.info(f"  - Log: training.log")
+    logger.info(f"  - Info: model_info.yaml")
+    if args.evaluate_rollout:
+        logger.info(f"  - Rollout evaluation: rollout_evaluation/*.png")
+
+
+if __name__ == '__main__':
+    main()
diff --git a/collab_env/gnn/interaction_particles/train.py b/collab_env/gnn/interaction_particles/train.py
new file mode 100644
index 0000000..92b4d34
--- /dev/null
+++ b/collab_env/gnn/interaction_particles/train.py
@@ -0,0 +1,1159 @@
+"""
+Training script for InteractionParticle model on 2D boids data.
+"""
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from torch_geometric.data import Data, Batch
+from torch_geometric.loader import DataLoader
+from tqdm import tqdm
+from loguru import logger
+import os
+import gc
+
+from .model import InteractionParticle
+from collab_env.gnn.utility import handle_discrete_data
+from collab_env.data.file_utils import expand_path
+
+
+def build_graph_from_data(positions, velocities, accelerations, next_positions=None, visual_range=0.5, max_radius=1.0):
+    """
+    Build PyG graph from position, velocity, and acceleration data.
+
+    Parameters
+    ----------
+    positions : torch.Tensor
+        Positions [N, 2]
+    velocities : torch.Tensor
+        Velocities [N, 2]
+    accelerations : torch.Tensor
+        Accelerations [N, 2] (targets for acceleration-based loss)
+    next_positions : torch.Tensor, optional
+        Next positions [N, 2] (targets for position-based loss)
+    visual_range : float
+        Maximum distance for edges (normalized)
+    max_radius : float
+        Maximum radius for normalization
+
+    Returns
+    -------
+    data : torch_geometric.data.Data
+        Graph data
+    """
+    N = positions.shape[0]
+
+    # Build edges based on visual range
+    dist = torch.cdist(positions, positions, p=2)
+    adj = (dist < visual_range).to(torch.bool)
+    adj.fill_diagonal_(False)
+    edge_index = adj.nonzero(as_tuple=False).T
+
+    # Create node features: [particle_id, pos_x, pos_y, vel_x, vel_y]
+    particle_ids = torch.arange(N, dtype=torch.float32).unsqueeze(1)
+    node_features = torch.cat([particle_ids, positions, velocities], dim=1)
+
+    # Create graph
+    data = Data(
+        x=node_features,
+        edge_index=edge_index,
+        y=accelerations,
+        pos=positions,
+        next_pos=next_positions if next_positions is not None else positions  # Store next position for position-based loss
+    )
+
+    return data
+
+
+def compute_velocity_statistics(dataset_path):
+    """
+    Compute velocity statistics from dataset for visualization purposes.
+
+    Parameters
+    ----------
+    dataset_path : str
+        Path to AnimalTrajectoryDataset .pt file
+
+    Returns
+    -------
+    dict with keys:
+        - 'max_speed': Maximum speed (velocity magnitude) in dataset
+        - 'mean_speed': Mean speed
+        - 'max_relative_velocity': Maximum relative velocity magnitude between particles
+        - 'mean_relative_velocity': Mean relative velocity magnitude
+    """
+    logger.info(f"Computing velocity statistics from {dataset_path}")
+
+    # Load dataset
+    dataset = torch.load(dataset_path, weights_only=False)
+
+    # Extract all positions
+    all_positions = []
+    for i in range(len(dataset)):
+        pos, species = dataset[i]
+        all_positions.append(pos.unsqueeze(0))
+
+    position = torch.cat(all_positions, dim=0)  # [B, T, N, 2]
+
+    # Compute velocities
+    velocities = torch.diff(position, dim=1)  # [B, T-1, N, 2]
+
+    # Compute speeds
+    speeds = torch.sqrt(torch.sum(velocities ** 2, dim=-1))  # [B, T-1, N]
+    max_speed = speeds.max().item()
+    mean_speed = speeds.mean().item()
+
+    # Compute pairwise relative velocities
+    B, T, N, D = velocities.shape
+    max_rel_vel = 0.0
+    sum_rel_vel = 0.0
+    count = 0
+
+    for b in range(min(B, 10)):  # Sample first 10 trajectories for efficiency
+        for t in range(T):
+            vel_t = velocities[b, t]  # [N, 2]
+            # Compute all pairwise relative velocities
+            rel_vel = vel_t.unsqueeze(0) - vel_t.unsqueeze(1)  # [N, N, 2]
+            rel_vel_mag = torch.sqrt(torch.sum(rel_vel ** 2, dim=-1))  # [N, N]
+            # Exclude self (diagonal)
+            mask = ~torch.eye(N, dtype=torch.bool)
+            rel_vel_mag_filtered = rel_vel_mag[mask]
+
+            max_rel_vel = max(max_rel_vel, rel_vel_mag_filtered.max().item())
+            sum_rel_vel += rel_vel_mag_filtered.sum().item()
+            count += len(rel_vel_mag_filtered)
+
+    mean_rel_vel = sum_rel_vel / count if count > 0 else 0.0
+
+    stats = {
+        'max_speed': max_speed,
+        'mean_speed': mean_speed,
+        'max_relative_velocity': max_rel_vel,
+        'mean_relative_velocity': mean_rel_vel
+    }
+
+    logger.info(f"Velocity statistics:")
+    logger.info(f"  Max speed: {max_speed:.4f}")
+    logger.info(f"  Mean speed: {mean_speed:.4f}")
+    logger.info(f"  Max relative velocity: {max_rel_vel:.4f}")
+    logger.info(f"  Mean relative velocity: {mean_rel_vel:.4f}")
+
+    return stats
+
+
+def prepare_dataset(dataset_path, visual_range=0.5, max_radius=1.0, input_differentiation='finite_diff', trajectory_indices=None):
+    """
+    Load and prepare 2D boids dataset for training.
+
+    Parameters
+    ----------
+    dataset_path : str
+        Path to AnimalTrajectoryDataset .pt file
+    visual_range : float
+        Maximum distance for edges (in normalized coordinates)
+    max_radius : float
+        Maximum radius for normalization
+    input_differentiation : str
+        Method for computing velocities/accelerations ('finite_diff' or 'spline')
+    trajectory_indices : list of int, optional
+        If provided, only process these trajectory indices from the dataset
+
+    Returns
+    -------
+    graphs : list
+        List of PyG Data objects
+    p_range : float
+        Position range for normalization
+
+    Notes
+    -----
+    The simulation has a natural timestep of 1.0 between frames.
+    Velocities and accelerations are computed as-is without rescaling.
+    """
+    logger.info(f"Loading 2D boids dataset from {dataset_path}")
+
+    # Load AnimalTrajectoryDataset
+    dataset = torch.load(dataset_path, weights_only=False)
+    logger.info(f"Loaded dataset with {len(dataset)} samples")
+
+    # Determine which trajectories to process
+    if trajectory_indices is not None:
+        indices_to_process = trajectory_indices
+        logger.info(f"Processing {len(indices_to_process)} selected trajectories")
+    else:
+        indices_to_process = range(len(dataset))
+        logger.info(f"Processing all {len(dataset)} trajectories")
+
+    # Extract positions from selected trajectories
+    # Each sample is (positions, species) where positions: [steps, N, 2]
+    all_positions = []
+    for i in indices_to_process:
+        pos, species = dataset[i]  # [steps, N, 2], [N]
+        # Add batch dimension: [1, steps, N, 2]
+        all_positions.append(pos.unsqueeze(0))
+
+    # Stack all samples: [B, steps, N, 2]
+    position = torch.cat(all_positions, dim=0)
+    logger.info(f"Extracted positions with shape: {position.shape}")
+
+    # Note: 2D boids data is already normalized by scene size [width, height]
+    # Positions are in [0, 1] range
+    # We just need to compute velocities and accelerations
+
+    # Compute velocities and accelerations using finite differences
+    # Natural timestep is 1.0 between frames
+    # v[t] = p[t+1] - p[t] (displacement per timestep)
+    # a[t] = v[t+1] - v[t] (velocity change per timestep)
+    p_smooth, v_smooth, a_smooth, v_function = handle_discrete_data(
+        position, input_differentiation
+    )
+
+    # Use as-is (no rescaling)
+    p_normalized = p_smooth.float()
+    v_normalized = v_smooth.float()
+    a_normalized = a_smooth.float()
+
+    # For 2D boids, data is already normalized
+    # p_range is 1.0 (scene is normalized to [0,1] x [0,1])
+    p_range = 1.0
+
+    # Build graphs
+    B, F, N, D = p_normalized.shape
+    logger.info(f"Building graphs from {B} trajectories, {F} frames, {N} particles")
+
+    graphs = []
+    for b in range(B):
+        for f in range(1, F-1):  # Skip first and last frame (no acceleration)
+            if torch.isnan(a_normalized[b, f]).any() or torch.isinf(a_normalized[b, f]).any():
+                continue
+
+            # Get next position for position-based loss
+            next_pos = p_normalized[b, f+1] if f+1 < F else p_normalized[b, f]
+
+            graph = build_graph_from_data(
+                p_normalized[b, f],
+                v_normalized[b, f],
+                a_normalized[b, f],
+                next_positions=next_pos,
+                visual_range=visual_range,
+                max_radius=max_radius
+            )
+            graphs.append(graph)
+
+    logger.info(f"Created {len(graphs)} graph samples")
+    return graphs, p_range
+
+
+def create_multistep_dataset(dataset_path, rollout_steps, visual_range=0.5, max_radius=1.0,
+                             input_differentiation='finite_diff'):
+    """
+    Create dataset of multi-step trajectory chunks for rollout training.
+
+    Parameters
+    ----------
+    dataset_path : str
+        Path to AnimalTrajectoryDataset .pt file
+    rollout_steps : int
+        Length of trajectory chunks (M)
+    visual_range : float
+        Maximum distance for edges (in normalized coordinates)
+    max_radius : float
+        Maximum radius for normalization
+    input_differentiation : str
+        Method for computing velocities/accelerations ('finite_diff' or 'spline')
+
+    Returns
+    -------
+    sequences : list of dict
+        Each dict contains:
+        - 'initial_pos': torch.Tensor [N, 2]
+        - 'initial_vel': torch.Tensor [N, 2]
+        - 'gt_positions': torch.Tensor [rollout_steps, N, 2]
+        - 'gt_velocities': torch.Tensor [rollout_steps, N, 2]
+    p_range : float
+        Position range for normalization
+
+    Notes
+    -----
+    The simulation has a natural timestep of 1.0 between frames.
+    Velocities are computed as-is without rescaling.
+    """
+    logger.info(f"Creating multi-step dataset with rollout_steps={rollout_steps}")
+    logger.info(f"Loading 2D boids dataset from {dataset_path}")
+
+    # Load AnimalTrajectoryDataset
+    dataset = torch.load(dataset_path, weights_only=False)
+    logger.info(f"Loaded dataset with {len(dataset)} samples")
+
+    # Extract all positions from the dataset
+    all_positions = []
+    for i in range(len(dataset)):
+        pos, species = dataset[i]  # [steps, N, 2], [N]
+        all_positions.append(pos.unsqueeze(0))
+
+    # Stack all samples: [B, steps, N, 2]
+    position = torch.cat(all_positions, dim=0)
+    logger.info(f"Extracted positions with shape: {position.shape}")
+
+    # Compute velocities (natural timestep = 1.0)
+    p_smooth, v_smooth, a_smooth, v_function = handle_discrete_data(
+        position, input_differentiation
+    )
+
+    # Use as-is (no rescaling)
+    v_normalized = v_smooth.float()
+    p_normalized = p_smooth.float()
+
+    p_range = 1.0
+
+    # Build trajectory sequences
+    B, F, N, D = p_normalized.shape
+    logger.info(f"Building trajectory sequences from {B} trajectories, {F} frames, {N} particles")
+    logger.info(f"Rollout steps: {rollout_steps}, need at least {rollout_steps} frames per trajectory")
+
+    if F < rollout_steps:
+        raise ValueError(
+            f"Dataset trajectories have only {F} frames, but rollout_steps={rollout_steps}. "
+            f"Either reduce --rollout-steps to <{F}, or use a dataset with longer trajectories."
+        )
+
+    sequences = []
+    nan_count = 0
+    for b in range(B):
+        # Use NON-OVERLAPPING windows to avoid data leakage
+        # Step by rollout_steps to get independent chunks
+        for start_frame in range(0, F - rollout_steps + 1, rollout_steps):
+            # Extract chunk of rollout_steps frames
+            pos_chunk = p_normalized[b, start_frame:start_frame + rollout_steps]  # [M, N, 2]
+            vel_chunk = v_normalized[b, start_frame:start_frame + rollout_steps]  # [M, N, 2]
+
+            # Check for NaN/Inf
+            if torch.isnan(pos_chunk).any() or torch.isinf(pos_chunk).any():
+                nan_count += 1
+                continue
+            if torch.isnan(vel_chunk).any() or torch.isinf(vel_chunk).any():
+                nan_count += 1
+                continue
+
+            sequences.append({
+                'initial_pos': pos_chunk[0].clone(),  # [N, 2]
+                'initial_vel': vel_chunk[0].clone(),  # [N, 2]
+                'gt_positions': pos_chunk.clone(),    # [M, N, 2]
+                'gt_velocities': vel_chunk.clone(),   # [M, N, 2]
+            })
+
+    logger.info(f"Created {len(sequences)} non-overlapping trajectory sequences (skipped {nan_count} due to NaN/Inf)")
+    return sequences, p_range
+
+
+def train_interaction_particle(
+    dataset_path,
+    config=None,
+    epochs=100,
+    batch_size=32,
+    learning_rate=1e-5,
+    train_split=0.8,
+    visual_range=0.5,
+    rollout_steps=1,
+    device=None,
+    save_dir=None,
+    seed=42,
+    epoch_callback=None
+):
+    """
+    Train InteractionParticle model on 2D boids data.
+
+    Supports both single-step (rollout_steps=1) and multi-step rollout training.
+
+    Parameters
+    ----------
+    dataset_path : str
+        Path to dataset .pt file
+    config : dict, optional
+        Model configuration
+    epochs : int
+        Number of training epochs
+    batch_size : int
+        Batch size
+    learning_rate : float
+        Learning rate
+    train_split : float
+        Fraction of data for training
+    visual_range : float
+        Maximum distance for edges
+    rollout_steps : int
+        Number of rollout steps for training (default: 1 for single-step)
+        If > 1, uses multi-step rollout training with position-based loss
+    device : str or torch.device, optional
+        Device to use
+    save_dir : str, optional
+        Directory to save model checkpoints
+    seed : int
+        Random seed
+    epoch_callback : callable, optional
+        Callback function called after each validation epoch with signature:
+        callback(model, epoch, train_loss, val_loss, save_dir)
+
+    Returns
+    -------
+    model : InteractionParticle
+        Trained model
+    history : dict
+        Training history
+
+    Notes
+    -----
+    The simulation uses a natural timestep of 1.0 between frames.
+    For integration during rollout, dt=1.0 is used implicitly.
+    """
+    # Set random seed
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+    # Set device
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    logger.info(f"Using device: {device}")
+
+    # Log training mode
+    if rollout_steps > 1:
+        logger.info(f"Using multi-step rollout training with {rollout_steps} steps")
+        logger.info(f"Loss will be computed on POSITIONS (not accelerations)")
+    else:
+        logger.info("Using single-step training (acceleration prediction)")
+
+    logger.info("Using natural timestep dt=1.0 from simulation")
+
+    # Default config
+    if config is None:
+        config = {
+            'n_particles': 20,
+            'n_particle_types': 1,
+            'max_radius': 1.0,
+            'hidden_dim': 128,
+            'embedding_dim': 16,
+            'n_mp_layers': 3,
+            'input_size': 7,  # delta_pos(2) + r(1) + delta_vel(2) + pos_i(2)
+            'output_size': 2,
+        }
+
+    # Prepare dataset (different for single-step vs multi-step)
+    if rollout_steps == 1:
+        # Single-step: use individual frames
+        graphs, p_range = prepare_dataset(
+            dataset_path,
+            visual_range=visual_range,
+            max_radius=config['max_radius']
+        )
+
+        # Split into train and validation
+        n_train = int(len(graphs) * train_split)
+        train_data = graphs[:n_train]
+        val_data = graphs[n_train:]
+
+        logger.info(f"Train samples: {len(train_data)}, Val samples: {len(val_data)}")
+
+        # Create data loaders
+        train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
+        val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False)
+    else:
+        # Multi-step: use trajectory sequences
+        sequences, p_range = create_multistep_dataset(
+            dataset_path,
+            rollout_steps=rollout_steps,
+            visual_range=visual_range,
+            max_radius=config['max_radius']
+        )
+
+        # Split into train and validation
+        n_train = int(len(sequences) * train_split)
+        train_data = sequences[:n_train]
+        val_data = sequences[n_train:]
+
+        logger.info(f"Train sequences: {len(train_data)}, Val sequences: {len(val_data)}")
+
+        # Check if we have enough data
+        if len(train_data) == 0:
+            raise ValueError(
+                f"No training sequences created! Dataset may be too short for rollout_steps={rollout_steps}. "
+                f"Try reducing --rollout-steps or using a longer dataset."
+            )
+
+        # For multi-step, we'll manually batch in the training loop
+        # (can't use DataLoader directly with dict format)
+        train_loader = train_data
+        val_loader = val_data
+
+    # Create model
+    model = InteractionParticle(config, device=device)
+    model.to(device)
+
+    # Optimizer
+    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+
+    # Training history
+    history = {
+        'train_loss': [],
+        'val_loss': [],
+        'config': config,
+        'p_range': p_range,
+        'rollout_steps': rollout_steps
+    }
+
+    # Training loop
+    logger.info("Starting training...")
+    best_val_loss = float('inf')
+
+    for epoch in range(epochs):
+        # Training
+        model.train()
+        train_loss = 0.0
+        train_batches = 0
+
+        if rollout_steps == 1:
+            # Single-step training: position-based loss (one-step rollout)
+            for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs} [Train]", leave=False):
+                batch = batch.to(device)
+
+                optimizer.zero_grad()
+
+                # Forward pass: predict acceleration
+                pred_acc = model(batch)
+
+                # Perform one-step integration to get predicted next position
+                # Extract positions and velocities from batch
+                # batch.pos is [N, 2], velocities are in batch.x
+                current_pos = batch.pos
+                current_vel = batch.x[:, 3:5]  # Extract velocities from node features (columns 3:5)
+
+                # Semi-implicit Euler integration
+                next_vel = current_vel + pred_acc * 1.0  # dt = 1.0
+                pred_next_pos = current_pos + next_vel * 1.0
+
+                # Compute loss (position MSE)
+                loss = F.mse_loss(pred_next_pos, batch.next_pos)
+
+                # Backward pass
+                loss.backward()
+                optimizer.step()
+
+                train_loss += loss.item()
+                train_batches += 1
+        else:
+            # Multi-step rollout training
+            # Manually batch the sequences
+            import random
+            random.shuffle(train_loader)  # train_loader is just a list of sequences
+
+            for i in tqdm(range(0, len(train_loader), batch_size), desc=f"Epoch {epoch+1}/{epochs} [Train]", leave=False):
+                batch_sequences = train_loader[i:i+batch_size]
+
+                optimizer.zero_grad()
+
+                # Stack all initial conditions and ground truths
+                initial_pos_batch = torch.stack([seq['initial_pos'] for seq in batch_sequences]).to(device)
+                initial_vel_batch = torch.stack([seq['initial_vel'] for seq in batch_sequences]).to(device)
+                gt_positions_batch = torch.stack([seq['gt_positions'] for seq in batch_sequences]).to(device)
+
+                # Perform batched differentiable rollout (processes all sequences in parallel!)
+                pred_positions_batch, _ = differentiable_rollout_batched(
+                    model, initial_pos_batch, initial_vel_batch, rollout_steps,
+                    visual_range=visual_range, delta_t=1.0, device=device,
+                    requires_grad=True
+                )
+
+                # Compute batched loss
+                loss = F.mse_loss(pred_positions_batch, gt_positions_batch)
+
+                # Single backward pass
+                loss.backward()
+
+                # Update parameters
+                optimizer.step()
+
+                batch_loss = loss.item()
+
+                train_loss += batch_loss
+                train_batches += 1
+
+                # Periodically clear GPU/CPU cache to prevent memory fragmentation
+                if train_batches % 50 == 0:
+                    gc.collect()
+                    if device.type == 'cuda':
+                        torch.cuda.empty_cache()
+
+        # Average training loss
+        if train_batches > 0:
+            train_loss /= train_batches
+        else:
+            logger.warning("No training batches processed in this epoch!")
+            train_loss = float('inf')
+
+        # Validation
+        model.eval()
+        val_loss = 0.0
+        val_batches = 0
+
+        if rollout_steps == 1:
+            # Single-step validation: position-based loss (one-step rollout)
+            with torch.no_grad():
+                for batch in tqdm(val_loader, desc=f"Epoch {epoch+1}/{epochs} [Val]", leave=False):
+                    batch = batch.to(device)
+
+                    # Predict acceleration
+                    pred_acc = model(batch)
+
+                    # Perform one-step integration
+                    current_pos = batch.pos
+                    current_vel = batch.x[:, 3:5]
+
+                    # Semi-implicit Euler integration
+                    next_vel = current_vel + pred_acc * 1.0
+                    pred_next_pos = current_pos + next_vel * 1.0
+
+                    # Compute position loss
+                    loss = F.mse_loss(pred_next_pos, batch.next_pos)
+                    val_loss += loss.item()
+                    val_batches += 1
+        else:
+            # Multi-step rollout validation
+            with torch.no_grad():
+                for i in tqdm(range(0, len(val_loader), batch_size), desc=f"Epoch {epoch+1}/{epochs} [Val]", leave=False):
+                    batch_sequences = val_loader[i:i+batch_size]
+
+                    # Stack all initial conditions and ground truths
+                    initial_pos_batch = torch.stack([seq['initial_pos'] for seq in batch_sequences]).to(device)
+                    initial_vel_batch = torch.stack([seq['initial_vel'] for seq in batch_sequences]).to(device)
+                    gt_positions_batch = torch.stack([seq['gt_positions'] for seq in batch_sequences]).to(device)
+
+                    # Perform batched rollout
+                    pred_positions_batch, _ = differentiable_rollout_batched(
+                        model, initial_pos_batch, initial_vel_batch, rollout_steps,
+                        visual_range=visual_range, delta_t=1.0, device=device,
+                        requires_grad=False
+                    )
+
+                    # Compute batched loss
+                    loss = F.mse_loss(pred_positions_batch, gt_positions_batch)
+
+                    val_loss += loss.item()
+                    val_batches += 1
+
+        # Average validation loss
+        if val_batches > 0:
+            val_loss /= val_batches
+        else:
+            logger.warning("No validation batches processed in this epoch!")
+            val_loss = float('inf')
+
+        # Record history
+        history['train_loss'].append(train_loss)
+        history['val_loss'].append(val_loss)
+
+        # Log progress
+        logger.info(
+            f"Epoch {epoch+1}/{epochs}: "
+            f"Train Loss = {train_loss:.2g}, Val Loss = {val_loss:.2g}"
+        )
+
+        # Save best model
+        if save_dir and val_loss < best_val_loss:
+            best_val_loss = val_loss
+            os.makedirs(save_dir, exist_ok=True)
+            save_path = os.path.join(save_dir, 'best_model.pt')
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'config': config,
+                'train_loss': train_loss,
+                'val_loss': val_loss,
+                'p_range': p_range,
+                'rollout_steps': rollout_steps
+            }, save_path)
+            logger.info(f"Saved best model to {save_path}")
+
+        # Call epoch callback if provided
+        if epoch_callback is not None:
+            epoch_callback(model, epoch + 1, train_loss, val_loss, save_dir)
+
+        # Clear memory after each epoch
+        gc.collect()
+        if device.type == 'cuda':
+            torch.cuda.empty_cache()
+
+    # Save final model
+    if save_dir:
+        final_path = os.path.join(save_dir, 'final_model.pt')
+        torch.save({
+            'epoch': epochs,
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'config': config,
+            'history': history
+        }, final_path)
+        logger.info(f"Saved final model to {final_path}")
+
+    return model, history
+
+
+def evaluate_model(model, dataset_path, visual_range=0.5, device=None, train_split=None):
+    """
+    Evaluate model on dataset.
+
+    Parameters
+    ----------
+    model : InteractionParticle
+        Trained model
+    dataset_path : str
+        Path to dataset
+    visual_range : float
+        Maximum distance for edges
+    device : str or torch.device, optional
+        Device to use
+    train_split : float, optional
+        If provided, only evaluate on validation set (indices >= train_split * dataset_size)
+        If None, evaluates on entire dataset
+
+    Returns
+    -------
+    metrics : dict
+        Evaluation metrics
+    """
+    if device is None:
+        device = next(model.parameters()).device
+
+    # Determine which trajectories to evaluate
+    trajectory_indices = None
+    if train_split is not None:
+        # Load dataset to get size
+        dataset = torch.load(dataset_path, weights_only=False)
+        n_train = int(len(dataset) * train_split)
+        trajectory_indices = list(range(n_train, len(dataset)))
+        logger.info(f"Evaluating on validation set only ({len(trajectory_indices)} trajectories)")
+
+    # Prepare dataset (only validation set if train_split specified)
+    graphs, _ = prepare_dataset(dataset_path, visual_range=visual_range, trajectory_indices=trajectory_indices)
+    loader = DataLoader(graphs, batch_size=32, shuffle=False)
+
+    # Evaluate
+    model.eval()
+    total_loss = 0.0
+    total_samples = 0
+
+    with torch.no_grad():
+        for batch in tqdm(loader, desc="Evaluating"):
+            batch = batch.to(device)
+            pred = model(batch)
+            loss = F.mse_loss(pred, batch.y, reduction='sum')
+            total_loss += loss.item()
+            total_samples += batch.y.shape[0]
+
+    mse = total_loss / total_samples
+    rmse = np.sqrt(mse)
+
+    metrics = {
+        'mse': mse,
+        'rmse': rmse,
+        'n_samples': total_samples
+    }
+
+    logger.info(f"Evaluation: MSE = {mse:.6f}, RMSE = {rmse:.6f}")
+
+    return metrics
+
+
+def build_graph_from_current_state(positions, velocities, visual_range, device):
+    """
+    Build PyG graph from current positions and velocities (for rollout).
+
+    Parameters
+    ----------
+    positions : torch.Tensor
+        Positions [N, 2]
+    velocities : torch.Tensor
+        Velocities [N, 2]
+    visual_range : float
+        Maximum distance for edges
+    device : torch.device
+        Device to use
+
+    Returns
+    -------
+    data : torch_geometric.data.Data
+        Graph data
+    """
+    N = positions.shape[0]
+
+    # Build edges based on visual range
+    dist = torch.cdist(positions.unsqueeze(0), positions.unsqueeze(0), p=2).squeeze(0)
+    adj = (dist < visual_range).to(torch.bool)
+    adj.fill_diagonal_(False)
+    edge_index = adj.nonzero(as_tuple=False).T.to(device)
+
+    # Create node features: [particle_id, pos_x, pos_y, vel_x, vel_y]
+    particle_ids = torch.arange(N, dtype=torch.float32, device=device).unsqueeze(1)
+    node_features = torch.cat([particle_ids, positions, velocities], dim=1)
+
+    # Create dummy accelerations (not used for prediction)
+    dummy_acc = torch.zeros_like(positions, device=device)
+
+    # Create graph
+    data = Data(
+        x=node_features,
+        edge_index=edge_index,
+        y=dummy_acc,
+        pos=positions
+    ).to(device)
+
+    return data
+
+
+def differentiable_rollout_batched(model, initial_positions_batch, initial_velocities_batch, n_steps,
+                                   visual_range=0.5, delta_t=0.1, device=None, requires_grad=False):
+    """
+    Perform batched differentiable multi-step rollout for multiple sequences in parallel.
+
+    Parameters
+    ----------
+    model : InteractionParticle
+        Model to use for prediction
+    initial_positions_batch : torch.Tensor
+        Initial positions [B, N, 2] where B is batch size
+    initial_velocities_batch : torch.Tensor
+        Initial velocities [B, N, 2]
+    n_steps : int
+        Number of steps to predict (including initial state)
+    visual_range : float
+        Maximum distance for edges
+    delta_t : float
+        Time step for integration
+    device : torch.device, optional
+        Device to use
+    requires_grad : bool
+        If True, enables gradient computation
+
+    Returns
+    -------
+    predicted_positions : torch.Tensor
+        Predicted positions [B, n_steps, N, 2]
+    predicted_velocities : torch.Tensor
+        Predicted velocities [B, n_steps, N, 2]
+    """
+    if device is None:
+        device = next(model.parameters()).device
+
+    B, N, D = initial_positions_batch.shape
+
+    # Initialize
+    pos_batch = initial_positions_batch.clone().to(device)  # [B, N, 2]
+    vel_batch = initial_velocities_batch.clone().to(device)  # [B, N, 2]
+
+    # Store all predictions
+    all_positions = [pos_batch.unsqueeze(1)]  # [B, 1, N, 2]
+    all_velocities = [vel_batch.unsqueeze(1)]  # [B, 1, N, 2]
+
+    def rollout_step():
+        nonlocal pos_batch, vel_batch
+
+        for step in range(n_steps - 1):
+            # Build graphs for all sequences in the batch
+            graphs = []
+            for b in range(B):
+                graph = build_graph_from_current_state(
+                    pos_batch[b], vel_batch[b], visual_range, device
+                )
+                graphs.append(graph)
+
+            # Batch the graphs using PyG Batch
+            batched_graph = Batch.from_data_list(graphs)
+
+            # Single forward pass for all graphs in parallel!
+            pred_acc_batched = model(batched_graph)
+
+            # Reshape back to [B, N, 2]
+            pred_acc_batch = pred_acc_batched.view(B, N, D)
+
+            # Semi-implicit Euler integration for all sequences
+            vel_batch = vel_batch + pred_acc_batch * delta_t
+            pos_batch = pos_batch + vel_batch * delta_t
+
+            # Store predictions
+            all_positions.append(pos_batch.unsqueeze(1))  # [B, 1, N, 2]
+            all_velocities.append(vel_batch.unsqueeze(1))  # [B, 1, N, 2]
+
+    if requires_grad:
+        rollout_step()
+    else:
+        with torch.no_grad():
+            rollout_step()
+
+    # Concatenate: [B, n_steps, N, 2]
+    predicted_positions = torch.cat(all_positions, dim=1)
+    predicted_velocities = torch.cat(all_velocities, dim=1)
+
+    return predicted_positions, predicted_velocities
+
+
+def differentiable_rollout(model, initial_positions, initial_velocities, n_steps,
+                          visual_range=0.5, delta_t=0.1, device=None,
+                          requires_grad=False, return_intermediates=False):
+    """
+    Perform differentiable multi-step rollout.
+
+    This function can be used for both training (with gradients) and evaluation (without).
+    It performs Euler integration of the predicted accelerations.
+
+    Parameters
+    ----------
+    model : InteractionParticle
+        Model to use for prediction
+    initial_positions : torch.Tensor
+        Initial positions [N, 2]
+    initial_velocities : torch.Tensor
+        Initial velocities [N, 2]
+    n_steps : int
+        Number of steps to predict (including initial state)
+    visual_range : float
+        Maximum distance for edges
+    delta_t : float
+        Time step for integration (default: 0.1)
+    device : torch.device, optional
+        Device to use
+    requires_grad : bool
+        If True, enables gradient computation (for training)
+        If False, uses torch.no_grad() (for evaluation)
+    return_intermediates : bool
+        If True, returns all intermediate states on device
+        If False, returns final states moved to CPU
+
+    Returns
+    -------
+    predicted_positions : torch.Tensor
+        Predicted positions [n_steps, N, 2]
+    predicted_velocities : torch.Tensor
+        Predicted velocities [n_steps, N, 2]
+    """
+    if device is None:
+        device = next(model.parameters()).device
+
+    # Initialize
+    pos = initial_positions.clone().to(device)
+    vel = initial_velocities.clone().to(device)
+
+    predicted_positions = []
+    predicted_velocities = []
+
+    # Store initial state
+    if return_intermediates:
+        predicted_positions.append(pos)
+        predicted_velocities.append(vel)
+    else:
+        predicted_positions.append(pos.detach().cpu())
+        predicted_velocities.append(vel.detach().cpu())
+
+    # Choose gradient context
+    def rollout_step():
+        nonlocal pos, vel
+
+        for step in range(n_steps - 1):
+            # Build graph from current state
+            graph = build_graph_from_current_state(pos, vel, visual_range, device)
+
+            # Predict acceleration
+            pred_acc = model(graph)
+
+            # Semi-implicit Euler integration (matching ground truth in boid.py)
+            # Ground truth: velocity changes are applied, then position updated with NEW velocity
+            # Training data: a[t] = v[t+1] - v[t], so model predicts acceleration
+            # Correct integration: v_{t+1} = v_t + a_t*dt, then x_{t+1} = x_t + v_{t+1}*dt
+            vel = vel + pred_acc * delta_t  # Update velocity FIRST
+            pos = pos + vel * delta_t       # Update position with NEW velocity
+
+            # Store predictions
+            if return_intermediates:
+                predicted_positions.append(pos)
+                predicted_velocities.append(vel)
+            else:
+                predicted_positions.append(pos.detach().cpu())
+                predicted_velocities.append(vel.detach().cpu())
+
+    if requires_grad:
+        # Training mode: keep gradients
+        rollout_step()
+    else:
+        # Evaluation mode: no gradients
+        with torch.no_grad():
+            rollout_step()
+
+    predicted_positions = torch.stack(predicted_positions, dim=0)
+    predicted_velocities = torch.stack(predicted_velocities, dim=0)
+
+    return predicted_positions, predicted_velocities
+
+
+def generate_rollout(model, initial_positions, initial_velocities, n_steps,
+                    visual_range=0.5, delta_t=1.0, device=None):
+    """
+    Generate autoregressive rollout predictions (evaluation mode).
+
+    This is a convenience wrapper around differentiable_rollout for backward compatibility.
+
+    Parameters
+    ----------
+    model : InteractionParticle
+        Trained model
+    initial_positions : torch.Tensor
+        Initial positions [N, 2]
+    initial_velocities : torch.Tensor
+        Initial velocities [N, 2]
+    n_steps : int
+        Number of steps to predict
+    visual_range : float
+        Maximum distance for edges
+    delta_t : float
+        Time step for integration (default: 1.0, natural simulation timestep)
+    device : str or torch.device, optional
+        Device to use
+
+    Returns
+    -------
+    predicted_positions : torch.Tensor
+        Predicted positions [n_steps, N, 2]
+    predicted_velocities : torch.Tensor
+        Predicted velocities [n_steps, N, 2]
+    """
+    model.eval()
+    return differentiable_rollout(
+        model, initial_positions, initial_velocities, n_steps,
+        visual_range=visual_range, delta_t=delta_t, device=device,
+        requires_grad=False, return_intermediates=False
+    )
+
+
+def evaluate_rollout(model, dataset_path, visual_range=0.5, n_rollout_steps=50, delta_t=1.0, device=None, train_split=0.8):
+    """
+    Evaluate model with multi-step rollout on validation data.
+
+    Parameters
+    ----------
+    model : InteractionParticle
+        Trained model
+    dataset_path : str
+        Path to dataset
+    visual_range : float
+        Maximum distance for edges
+    n_rollout_steps : int
+        Number of steps for rollout
+    delta_t : float
+        Time step for integration (default: 1.0, natural simulation timestep)
+    device : str or torch.device, optional
+        Device to use
+    train_split : float
+        Train/val split ratio (default: 0.8) - should match training split
+
+    Returns
+    -------
+    results : dict
+        Dictionary containing:
+        - ground_truth_positions: list of [T, N, 2]
+        - predicted_positions: list of [T, N, 2]
+        - ground_truth_velocities: list of [T, N, 2]
+        - predicted_velocities: list of [T, N, 2]
+        - metrics: dict of aggregate metrics
+    """
+    if device is None:
+        device = next(model.parameters()).device
+
+    logger.info(f"Loading dataset for rollout evaluation from {dataset_path}")
+
+    # Load dataset
+    dataset = torch.load(dataset_path, weights_only=False)
+
+    # Use validation split (consistent with training)
+    n_train = int(len(dataset) * train_split)
+    val_dataset = [dataset[i] for i in range(n_train, len(dataset))]
+
+    logger.info(f"Using train_split={train_split:.2f}, evaluating rollout on {len(val_dataset)} validation samples")
+
+    all_gt_positions = []
+    all_pred_positions = []
+    all_gt_velocities = []
+    all_pred_velocities = []
+
+    position_errors = []
+    velocity_errors = []
+
+    model.eval()
+
+    for idx, (positions, species) in enumerate(tqdm(val_dataset, desc="Generating rollouts")):
+        # positions: [T, N, 2]
+        T, N, D = positions.shape
+
+        # Use available timesteps (need at least 2 for initial velocity)
+        if T < 2:
+            continue
+
+        # Limit rollout steps to available data
+        # generate_rollout returns n_steps timesteps (including initial)
+        actual_rollout_steps = min(n_rollout_steps, T)
+
+        # Extract ground truth trajectory (same number of timesteps as rollout will produce)
+        gt_positions = positions[:actual_rollout_steps].numpy()
+
+        # Compute ground truth velocities
+        gt_velocities = np.diff(gt_positions, axis=0)
+        gt_velocities = np.concatenate([gt_velocities, gt_velocities[-1:]], axis=0)  # Repeat last
+
+        # Get initial conditions
+        initial_pos = positions[0]  # [N, 2]
+        initial_vel = positions[1] - positions[0]  # [N, 2]
+
+        # Generate rollout
+        pred_pos, pred_vel = generate_rollout(
+            model, initial_pos, initial_vel,
+            n_steps=actual_rollout_steps,
+            visual_range=visual_range,
+            delta_t=delta_t,
+            device=device
+        )
+
+        # Convert to numpy
+        pred_pos = pred_pos.numpy()
+        pred_vel = pred_vel.numpy()
+
+        # Store results
+        all_gt_positions.append(gt_positions)
+        all_pred_positions.append(pred_pos)
+        all_gt_velocities.append(gt_velocities)
+        all_pred_velocities.append(pred_vel)
+
+        # Compute errors for this trajectory
+        pos_error = np.mean(np.linalg.norm(gt_positions - pred_pos, axis=-1))  # Mean over time and particles
+        vel_error = np.mean(np.linalg.norm(gt_velocities - pred_vel, axis=-1))
+
+        position_errors.append(pos_error)
+        velocity_errors.append(vel_error)
+
+    # Aggregate metrics
+    metrics = {
+        'mean_position_error': np.mean(position_errors),
+        'std_position_error': np.std(position_errors),
+        'mean_velocity_error': np.mean(velocity_errors),
+        'std_velocity_error': np.std(velocity_errors),
+        'n_trajectories': len(position_errors)
+    }
+
+    logger.info(f"Rollout evaluation complete:")
+    logger.info(f"  Mean position error: {metrics['mean_position_error']:.6f} ± {metrics['std_position_error']:.6f}")
+    logger.info(f"  Mean velocity error: {metrics['mean_velocity_error']:.6f} ± {metrics['std_velocity_error']:.6f}")
+
+    results = {
+        'ground_truth_positions': all_gt_positions,
+        'predicted_positions': all_pred_positions,
+        'ground_truth_velocities': all_gt_velocities,
+        'predicted_velocities': all_pred_velocities,
+        'metrics': metrics
+    }
+
+    return results
diff --git a/collab_env/gnn/interaction_particles/train_2d_boids.sh b/collab_env/gnn/interaction_particles/train_2d_boids.sh
new file mode 100755
index 0000000..bb224ff
--- /dev/null
+++ b/collab_env/gnn/interaction_particles/train_2d_boids.sh
@@ -0,0 +1,124 @@
+#!/bin/bash
+#
+# Quick start training script for InteractionParticle on 2D boids data
+#
+# This script provides example commands for training on existing 2D boids datasets.
+# Uncomment the command you want to run, or modify as needed.
+
+# Default: Quick test (10 epochs) on basic dataset
+# python -m collab_env.gnn.interaction_particles.run_training \
+#     --dataset simulated_data/boid_single_species_basic.pt \
+#     --epochs 10 \
+#     --batch-size 32 \
+#     --visual-range 0.104 \
+#     --save-dir trained_models/interaction_particle_2d_quick \
+#     --device auto
+
+# Full training (100 epochs) on basic dataset
+python -m collab_env.gnn.interaction_particles.run_training \
+    --dataset simulated_data/boid_single_species_basic.pt \
+    --epochs 2000 \
+    --batch-size 32 \
+    --visual-range 0.2 \
+    --plot-every 500 \
+    --evaluate-rollout \
+    --n-rollout-steps 10 \
+    --learning-rate 1e-3 \
+    --n-layers 3 \
+    --hidden-dim 64 \
+    --embedding-dim 0 \
+    --rollout-steps 10 \
+    --save-dir trained_models/interaction_particle_2d_basic_rollout10 \
+    --device cpu
+
+# weak alignment dataset
+python -m collab_env.gnn.interaction_particles.run_training \
+    --dataset simulated_data/boid_single_species_weakalignment_large.pt \
+    --epochs 50 \
+    --batch-size 32 \
+    --visual-range 0.5 \
+    --plot-every 2 \
+    --evaluate-rollout \
+    --n-rollout-steps 80 \
+    --learning-rate 1e-4 \
+    --n-layers 3 \
+    --hidden-dim 256 \
+    --embedding-dim 0 \
+    --rollout-steps 10 \
+    --save-dir trained_models/interaction_particle_2d_weakalignment_rollout10 \
+    --device cpu
+
+
+# Training on LARGE runpod dataset (1000 samples, T=800 timesteps)
+# NOTE: Requires much smaller batch size due to long trajectories
+python -m collab_env.gnn.interaction_particles.run_training \
+    --dataset simulated_data/runpod/boid_single_species_basic.pt \
+    --epochs 20 \
+    --batch-size 32 \
+    --visual-range 0.2 \
+    --plot-every 2 \
+    --evaluate-rollout \
+    --n-rollout-steps 50 \
+    --save-dir trained_models/interaction_particle_2d_runpod \
+    --device cpu \
+    --learning-rate 1e-4 \
+    --n-layers 4 \
+    --embedding-dim 0 \
+    --rollout-steps 10 \
+    --hidden-dim 64
+
+# weak alignment dataset
+python -m collab_env.gnn.interaction_particles.run_training \
+    --dataset simulated_data/runpod/boid_single_species_weakalignment_large.pt \
+    --epochs 5 \
+    --batch-size 32 \
+    --visual-range 0.2 \
+    --plot-every 1 \
+    --evaluate-rollout \
+    --n-rollout-steps 50 \
+    --save-dir trained_models/interaction_particle_2d_runpod_weakalignment \
+    --device cpu \
+    --learning-rate 1e-4 \
+    --n-layers 3 \
+    --embedding-dim 0 \
+    --hidden-dim 256
+
+# independent dataset
+python -m collab_env.gnn.interaction_particles.run_training \
+    --dataset simulated_data/runpod/boid_single_species_independent.pt \
+    --epochs 5 \
+    --batch-size 32 \
+    --visual-range 0.2 \
+    --plot-every 1 \
+    --evaluate-rollout \
+    --n-rollout-steps 50 \
+    --save-dir trained_models/interaction_particle_2d_runpod_independent \
+    --device cpu \
+    --learning-rate 1e-4 \
+    --n-layers 3 \
+    --embedding-dim 0 \
+    --hidden-dim 256
+
+# Use specific device (cpu, cuda, mps, or auto)
+# python -m collab_env.gnn.interaction_particles.run_training \
+#     --dataset simulated_data/boid_single_species_basic.pt \
+#     --epochs 100 \
+#     --batch-size 32 \
+#     --visual-range 0.104 \
+#     --device mps \
+#     --save-dir trained_models/interaction_particle_2d_mps
+
+# Notes:
+# - visual_range: 0.104 ≈ 50px / 480px (boid visual_range / scene_size)
+# - All datasets are in simulated_data/ directory
+# - Config files are auto-detected (e.g., *_config.pt)
+# - Results saved to trained_models/interaction_particle_2d_*/
+# - --device auto: Auto-detects cuda > mps > cpu
+# - --plot-every N: Generates force decomposition and rollout plots every N epochs
+#   Plots saved in epoch_XXXX/ subdirectories
+# - --evaluate-rollout: Enables multi-step rollout evaluation
+#   Results saved in rollout_evaluation/ subdirectory
+# - Final plots always generated:
+#   * training_history.png - loss curves
+#   * learned_force_decomposition.png - learned forces (2x3 grid)
+#   * true_boid_force_decomposition.png - ground truth forces (2x3 grid)
diff --git a/collab_env/gnn/utility.py b/collab_env/gnn/utility.py
index f8af15f..65d9a3d 100644
--- a/collab_env/gnn/utility.py
+++ b/collab_env/gnn/utility.py
@@ -77,15 +77,18 @@ def finite_diff(time, position, query=None):
 
 def finite_diff_data(position):
     """Like upgrade_data, but no smoothing, just upgrade_data_finite_diff"""
-    B, F, N, dim = np.shape(position)
-    # time = np.arange(F)
+    # Ensure position is a torch tensor
+    if not isinstance(position, torch.Tensor):
+        position = torch.from_numpy(position)
+
+    B, F, N, dim = position.shape
 
     v = torch.zeros_like(position)
     a = torch.zeros_like(position)
     # v_function = {}
 
-    v[:, 1:, :, :] = torch.diff(position, axis=1)
-    a[:, :-1, :, :] = torch.diff(v, axis=1)
+    v[:, 1:, :, :] = torch.diff(position, dim=1)
+    a[:, :-1, :, :] = torch.diff(v, dim=1)
 
     return position, v, a, None