Experimental Refactoring

.gitignore

@@ -14,6 +14,16 @@ __pycache__/
 *.pdf
 *.pt
 *.mp4
+results/
+
+# Data files
+*.csv
+*.jsonl
+*.json
+*.dat
+*.pkl
+*.hdf5
+*.h5
 *.json
 
 # Sketch

README.md

@@ -1,102 +1,149 @@
 # Diffusion Evolution
 
-This repo is for our ICLR 2025 paper [Diffusion models are evolutionary algorithms](https://openreview.net/forum?id=xVefsBbG2O), which anayatically proves that diffusion models are a type of evolutionary algorithm. This equivalence allows us to leverage advancements in diffusion models for evolutionary algorithm tasks, including accelerated sampling and latent space diffusion.
+This repository contains the official implementation of the ICLR 2025 paper, "[Diffusion Models are Evolutionary Algorithms](https://openreview.net/forum?id=xVefsBbG2O)". This work analytically proves that diffusion models can be interpreted as a form of evolutionary algorithm. This equivalence allows us to leverage advancements in diffusion models for evolutionary tasks, including accelerated sampling and latent space exploration.
 
-![](./experiments/2d_models/two_peaks/images/framwork.jpg)
+The core idea of the Diffusion Evolution framework is to treat the reverse diffusion process as an evolutionary algorithm. A population of samples estimates the noise that was added to them (or their noise-free states) based on savory of their neighbors. The population then "evolves" by taking a denoising step.
 
-The Diffusion Evolution framework treats inversed diffusion as evolutionary algorithm, where the population estimates its added noise (or their noise-free states) based on its neighbors' fitness then evolves via denoising. The following figure shows the process on optimizing a two-peak density function. The Diffusion Evolution initially has large neighbor range (shown as blue disk), calculating $x_0$ based on the fitness of its neighbors then move toward estimated $x_0$.
+## Project Structure
 
-![](./experiments/2d_models/figures/process.png)
+This project is organized with a focus on academic rigor, clarity, and reproducibility. The core library is located in `src/diffevo`, with a modular structure that separates concerns. The project now uses a plugin-based architecture for problems, optimizers, and callbacks.
 
+- `src/diffevo`: The core library, containing the `Orchestrator` and base classes for problems, optimizers, and callbacks.
+- `plugins/`: Home for problems, optimizers, and callbacks.
+  - `plugins/problems/`: Problem definitions.
+  - `plugins/optimizers/`: Optimizer implementations.
+  - `plugins/callbacks/`: Callback implementations.
+- `configs/`: YAML configuration files for experiments.
+- `tests/`: Unit tests for the core library.
 
-## Install
+## Installation
 
-You can install the package via pip:
+Getting started is as simple as cloning the repository and running the initialization script. This will handle the creation of a virtual environment and the installation of all necessary dependencies.
 
 ```bash
-pip install diffevo
+# Clone the repository
+git clone https://github.com/Zhangyanbo/diffusion-evolution
+cd diffusion-evolution
+
+# Run the initialization script
+python init.py
 ```
 
-or manually install:
+The `init.py` script will create a local Python environment in a `.venv` directory and install all dependencies. If you run it again, it will prompt you to reset the environment.
+
+## Quick Start
+
+To see the script in action, you can run a quick "smoketest" to verify that everything is working correctly.
+
 ```bash
-clone https://github.com/Zhangyanbo/diffusion-evolution
-cd diffevo/
-pip install .
+# Make sure to activate the virtual environment first
+source .venv/bin/activate
+
+# Run the smoketest
+python run.py configs/smoketest.yaml --smoketest
 ```
 
-Some benchmark codes requires dependencies, can be installed via:
+This will run a minimal configuration and save the results to a timestamped directory in `results/`.
+
+## Running Experiments
+
+All experiments are run through `run.py`, which takes a YAML configuration file as an argument.
+
 ```bash
-pip install cma gym pygame tqdm matplotlib numpy==1.26.4 
+python run.py <path_to_config>.yaml
 ```
 
-Also Pytorch version 2.5 or above is required
+We provide several example configurations in the `configs/` directory.
+
+## Generating a Report
+
+To run a sequence of experiments and generate a report, you can use the `report.py` script.
+
 ```bash
-pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
+python report.py
 ```
 
+This will run a predefined sequence of experiments (currently, just the smoketest) and will eventually generate a report of the results.
 
-The benchmark fitness functions can be found here: https://github.com/bhartl/foobench 
+## Graph-Based Optimizers
 
-## Typical Usage
+This framework includes two distinct approaches to using graphs in Diffusion Evolution:
 
-In most cases, tuning hyperparameters or adding custom operations is necessary to achieve higher performance. We recommend using the following form for the best balance between conciseness and versatility.
+1.  **`GraphDiffEvo`**: For problems *defined on* a graph.
+2.  **`GraphBasedDiffEvo`**: An optimizer where the population *is* a graph.
 
-```python
-from diffevo import DDIMScheduler, BayesianGenerator
-from diffevo.examples import two_peak_density
+### 1. `GraphDiffEvo`: Solving Problems on Graphs
 
-scheduler = DDIMScheduler(num_step=100)
+This is a specialized optimizer for graph-based optimization problems. The goal is to find a graph structure that maximizes a certain objective function.
 
-x = torch.randn(512, 2)
+#### Running `GraphDiffEvo` Experiments
 
-for t, alpha in scheduler:
-    fitness = two_peak_density(x, std=0.25)
-    generator = BayesianGenerator(x, fitness, alpha)
-    x = generator(noise=0)
+We provide two example graph-based experiments:
+
+-   **Graph Flow**: This experiment attempts to evolve a graph that maximizes the "flow" of capacity from a source node to a sink node. To run it:
+    ```bash
+    python run.py configs/graph_flow.yaml
+    ```
+
+-   **Max Clique**: This experiment attempts to find the largest clique (a fully connected subgraph) in a graph. To run it:
+    ```bash
+    python run.py configs/max_clique.yaml
+    ```
+
+### 2. `GraphBasedDiffEvo`: Structuring the Population as a Graph
+
+This optimizer introduces a novel approach where the population of candidate solutions is itself structured as a sparse graph. The evolutionary updates (diffusion) only occur between connected nodes (neighbors), rather than across the entire population. This can improve efficiency and exploration for large populations.
+
+#### Running `GraphBasedDiffEvo` Experiments
+
+To run an experiment with this optimizer, you can use the example configuration:
+```bash
+python run.py configs/graph_based_evolution.yaml
 ```
 
-The generator requires fitness values to be non-negative. If your objective
-returns negative values, please apply a mapping (see `diffevo.fitnessmapping`)
-to convert them before calling the generator.
+This will run the Rosenbrock benchmark problem, but the `GraphBasedDiffEvo` optimizer will use a k-Nearest Neighbors graph to structure its internal population.
 
-The following are two evolution trajectories of different fitness functions.
+This repository includes a specialized framework for graph-based optimization problems using Diffusion Evolution. This framework is designed to be extensible, allowing researchers to easily define new graph problems and apply the `GraphDiffEvo` optimizer to them.
 
-## Advanced Usage
+### Running Graph Evolution Experiments
 
-We also offer multiple choices for each component to accommodate more advanced use cases:
+We provide two example graph-based experiments:
 
-* In addition to the `DDIMScheduler`, we provide the `DDIMSchedulerCosine`, which features a different $\alpha$ scheduler.
-* We offer multiple fitness mapping functions that map the original fitness to a different value. These can be found in `diffevo.fitnessmapping`.
-* Currently, we have only one version of the generator.
+1.  **Graph Flow**: This experiment attempts to evolve a graph that maximizes the "flow" of capacity from a source node to a sink node. To run it:
+    ```bash
+    python run.py configs/graph_flow.yaml
+    ```
 
-Below is an example of how to change the diffusion process and conduct advanced experiments:
+2.  **Max Clique**: This experiment attempts to find the largest clique (a fully connected subgraph) in a graph. To run it:
+    ```bash
+    python run.py configs/max_clique.yaml
+    ```
 
-```python
-import torch
-from diffevo import DDIMScheduler, BayesianGenerator, DDIMSchedulerCosine
-from diffevo.examples import two_peak_density
-from diffevo.fitnessmapping import Power, Energy, Identity
+### Creating a New Plugin
 
-scheduler = DDIMSchedulerCosine(num_step=100) # use a different scheduler
+To create a new plugin, you need to:
 
-x = torch.randn(512, 2)
+1.  **Create a New Python File**: Create a new Python file in the appropriate plugin directory (`plugins/problems`, `plugins/optimizers`, or `plugins/callbacks`).
+2.  **Create a New Class**: In the new file, create a new class that inherits from the appropriate base class (`Problem`, `Optimizer`, or `Callback`).
+3.  **Implement the Required Methods**: Implement the required methods for the base class. For example, a `Problem` plugin needs to implement `evaluate`.
+4.  **Update Your Configuration File**: In your YAML configuration file, update the `problem.name`, `optimizer.class_name`, or `callbacks` list to match the name of your new class.
 
-trace = [] # store the trace of the population
+Here is a simple template for a new problem plugin:
 
-mapping_fn = Power(3) # setup the power mapping function
+```python
+# In plugins/problems/my_problem.py
+from diffevo.problems.base import Problem
+import torch
 
-for t, alpha in scheduler:
-    fitness = two_peak_density(x, std=0.25)
-    # apply the power mapping function
-    generator = BayesianGenerator(x, mapping_fn(fitness), alpha)
-    x = generator(noise=0.1)
-    trace.append(x)
+class MyProblem(Problem):
+    def __init__(self, dim):
+        super().__init__(dim=dim)
 
-trace = torch.stack(trace)
+    def evaluate(self, x):
+        return torch.sum(x ** 2, dim=-1)
 ```
 
-
-### Cite our work
+## Citing Our Work
 
 ```
 @inproceedings{
@@ -111,4 +158,4 @@ url={https://openreview.net/forum?id=xVefsBbG2O}
 
 ## License
 
-Our software is relased under modified Apache 2.0 License. We allow non-commercial usage for research, study, learning, etc., while limiting the commercial usage.
+This software is released under the Apache 2.0 License. See the [LICENSE](LICENSE) file for more details.

configs/base.yaml

@@ -0,0 +1,8 @@
+# Base configuration for all experiments
+seed: 42
+num_runs: 1
+
+callbacks:
+  - CSVLogger
+  - ConsoleLogger
+  - PlottingCallback

configs/graph_based_evolution.yaml

@@ -0,0 +1,16 @@
+base: base.yaml
+
+name: graph_based_evolution_smoketest
+
+optimizer:
+  module: diffevo.optimizers.graph_based_diffevo
+  class_name: GraphBasedDiffEvo
+  params:
+    num_step: 100
+    popsize: 50
+    k: 5
+
+problem:
+  name: Rosenbrock
+  params:
+    dim: 2

configs/graph_evolution.yaml

@@ -0,0 +1,20 @@
+# Configuration for the graph evolution experiment
+
+problem:
+  name: Graphflow
+  params:
+    num_nodes: 3
+
+optimizer:
+  name: GraphDiffEvo
+  params:
+    pop_size: 512
+    dim: 9
+    T: 100
+    sigma_m: 0.0
+    power: 3.0
+
+callbacks:
+  - name: CSVLogger
+
+n_runs: 1

configs/graph_flow.yaml

@@ -0,0 +1,17 @@
+base: base.yaml
+
+name: graph_flow_experiment
+
+problem:
+  name: Graphflow
+  params:
+    dim: 25
+    num_nodes: 5
+
+optimizer:
+  name: GraphDiffEvo
+  params:
+    pop_size: 512
+    num_steps: 100
+    sigma_m: 0.0
+    power: 3.0

configs/image_evolution.yaml

@@ -0,0 +1,18 @@
+# Configuration for the image evolution experiment
+
+problem:
+  name: Image
+  params:
+    dim_sqrt: 28
+    target_image_name: "mnist_7"
+
+optimizer:
+  name: DiffEvo
+  params:
+    pop_size: 64
+    max_iters: 100
+
+callbacks:
+  - name: CSVLogger
+
+n_runs: 1

configs/max_clique.yaml

@@ -0,0 +1,17 @@
+base: base.yaml
+
+name: max_clique_experiment
+
+problem:
+  name: Maxclique
+  params:
+    dim: 25
+    num_nodes: 5
+
+optimizer:
+  name: GraphDiffEvo
+  params:
+    pop_size: 512
+    num_steps: 100
+    sigma_m: 0.0
+    power: 3.0

configs/rl_diffevo.yaml

@@ -0,0 +1,17 @@
+base: smoketest.yaml
+
+problem:
+  name: RL
+  env_name: "CartPole-v1"
+  dim_hidden: 8
+  n_hidden_layers: 1
+
+optimizer:
+  name: RLEvo
+  num_step: 10
+  population_size: 256
+  T: 10
+  scaling: 100
+  latent_dim: null
+  noise: 1
+  weight_decay: 0

configs/rl_evolution.yaml

@@ -0,0 +1,17 @@
+# Configuration for the RL evolution experiment
+
+problem:
+  name: RL
+  params:
+    env_name: "CartPole-v1"
+
+optimizer:
+  name: DiffEvo
+  params:
+    pop_size: 64
+    max_iters: 100
+
+callbacks:
+  - name: CSVLogger
+
+n_runs: 1

configs/smoketest.yaml

@@ -0,0 +1,15 @@
+base: base.yaml
+
+name: smoketest
+
+optimizer:
+  module: diffevo.optimizers.diffevo
+  class_name: DiffEvo
+  params:
+    num_step: 1
+    popsize: 10
+
+problem:
+  name: Rosenbrock
+  params:
+    dim: 2