# `Edifice.Generative.Diffusion`
[🔗](https://github.com/blasphemetheus/edifice/blob/main/lib/edifice/generative/diffusion.ex#L1)

Diffusion Policy: Action generation via denoising diffusion.

Implements Diffusion Policy from "Diffusion Policy: Visuomotor Policy Learning
via Action Diffusion" (Chi et al., RSS 2023). Instead of directly predicting
actions, we learn to denoise random noise into actions conditioned on observations.

## Key Innovation: DDPM for Actions

Traditional policies: `a = pi(o)` - direct mapping
Diffusion Policy: `a = denoise(noise | o)` - iterative refinement

```
Training:
  1. Sample action sequence a_0 from data
  2. Add noise: a_t = sqrt(alpha_bar_t)*a_0 + sqrt(1-alpha_bar_t)*eps
  3. Predict noise: eps_hat = network(a_t, t, obs)
  4. Loss: ||eps - eps_hat||^2

Inference:
  1. Sample a_T ~ N(0, I)
  2. For t = T...1: a_{t-1} = denoise(a_t, t, obs)
  3. Return a_0
```

## Architecture

```
Observations [batch, obs_dim]
      |
      v
+-------------------------------------+
|  Observation Encoder                 |
|  (MLP or temporal backbone)          |
+-------------------------------------+
      |
      v obs_embed
+-------------------------------------+
|  Denoising Network                   |
|  Input: (noisy_actions, timestep,    |
|          obs_embed)                  |
|  Output: predicted_noise             |
+-------------------------------------+
      |
      v
Denoised Actions [batch, action_horizon, action_dim]
```

## Advantages

| Feature | Benefit |
|---------|---------|
| Multi-modal | Can represent multiple valid actions |
| High-dim | Scales well to action sequences |
| Stable | MSE loss is simple and stable |
| Expressive | Captures complex action distributions |

## Usage

    # Build diffusion policy
    model = Diffusion.build(
      obs_size: 287,
      action_dim: 64,
      action_horizon: 8,
      num_diffusion_steps: 100
    )

    # Training: predict noise
    {loss, predicted_noise} = Diffusion.training_step(
      model, params, observations, actions, key
    )

    # Inference: denoise to get actions
    actions = Diffusion.sample(model, params, observations, key)

## References
- Paper: https://arxiv.org/abs/2303.04137
- Project: https://diffusion-policy.cs.columbia.edu/

# `build_opt`

```elixir
@type build_opt() ::
  {:obs_size, pos_integer()}
  | {:action_dim, pos_integer()}
  | {:action_horizon, pos_integer()}
  | {:hidden_size, pos_integer()}
  | {:num_layers, pos_integer()}
  | {:num_steps, pos_integer()}
```

Options for `build/1`.

# `build`

```elixir
@spec build([build_opt()]) :: Axon.t()
```

Build a Diffusion Policy model.

## Options
  - `:obs_size` - Size of observation embedding (required)
  - `:action_dim` - Dimension of action space (required)
  - `:action_horizon` - Number of actions to predict (default: 8)
  - `:hidden_size` - Hidden dimension (default: 256)
  - `:num_layers` - Number of denoiser layers (default: 4)
  - `:num_steps` - Number of diffusion timesteps (default: 100)

## Returns
  An Axon model that predicts noise given (noisy_actions, timestep, obs).

# `build_denoiser`

```elixir
@spec build_denoiser(Axon.t(), Axon.t(), Axon.t(), keyword()) :: Axon.t()
```

Build the denoising network (noise predictor).

Architecture: MLP with sinusoidal timestep embedding and observation conditioning.

# `build_obs_encoder`

```elixir
@spec build_obs_encoder(keyword()) :: Axon.t()
```

Build the observation encoder for temporal inputs.

Processes sequence of observations into a single embedding.

# `compute_loss`

```elixir
@spec compute_loss(Nx.Tensor.t(), Nx.Tensor.t()) :: Nx.Tensor.t()
```

Compute training loss: MSE between true and predicted noise.

# `default_action_horizon`

```elixir
@spec default_action_horizon() :: pos_integer()
```

Default action prediction horizon

# `default_beta_end`

```elixir
@spec default_beta_end() :: float()
```

Beta schedule end

# `default_beta_start`

```elixir
@spec default_beta_start() :: float()
```

Beta schedule start

# `default_hidden_size`

```elixir
@spec default_hidden_size() :: pos_integer()
```

Default hidden dimension

# `default_num_layers`

```elixir
@spec default_num_layers() :: pos_integer()
```

Default number of denoiser layers

# `default_num_steps`

```elixir
@spec default_num_steps() :: pos_integer()
```

Default number of diffusion timesteps

# `fast_inference_defaults`

```elixir
@spec fast_inference_defaults() :: keyword()
```

Fast inference configuration with fewer diffusion steps.

# `make_schedule`

```elixir
@spec make_schedule(keyword()) :: map()
```

Precompute diffusion schedule constants.

Returns a map with:
- `:betas` - Noise schedule beta_t
- `:alphas` - 1 - beta_t
- `:alphas_cumprod` - alpha_bar_t = product of alpha_s
- `:sqrt_alphas_cumprod` - sqrt(alpha_bar_t)
- `:sqrt_one_minus_alphas_cumprod` - sqrt(1-alpha_bar_t)
- `:sqrt_recip_alphas` - 1/sqrt(alpha_t)
- `:posterior_variance` - beta_tilde_t for sampling

# `output_size`

```elixir
@spec output_size(keyword()) :: non_neg_integer()
```

Get the output size of a Diffusion Policy model.

Returns action_horizon * action_dim.

# `p_sample`

```elixir
@spec p_sample(Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), map()) ::
  Nx.Tensor.t()
```

Single denoising step (reverse process).

```
a_{t-1} = (1/sqrt(alpha_t)) * (a_t - (beta_t/sqrt(1-alpha_bar_t)) * eps_hat) + sqrt(beta_tilde_t) * z
```

# `param_count`

```elixir
@spec param_count(keyword()) :: non_neg_integer()
```

Calculate approximate parameter count for a Diffusion Policy model.

# `q_sample`

```elixir
@spec q_sample(Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), map()) :: Nx.Tensor.t()
```

Forward diffusion: add noise to actions.

```
a_t = sqrt(alpha_bar_t) * a_0 + sqrt(1-alpha_bar_t) * eps
```

# `recommended_defaults`

```elixir
@spec recommended_defaults() :: keyword()
```

Get recommended defaults for action generation.

---

*Consult [api-reference.md](api-reference.md) for complete listing*