# Top-K Discovery Metrics

`TopKMetricsModule` measures the quality and diversity of the highest-reward
samples drawn from the learned forward policy. It repeatedly samples trajectories,
picks the top `K` rewards, and reports aggregate statistics for each requested
`K`.

## Intuition

- Monitor whether the forward policy can surface high-reward terminals by
  sampling a large batch, exponentiating log rewards, and keeping the best `k`
  entries per evaluation run.
- The diversity statistic highlights mode collapse: a low value with high
  reward means the policy keeps finding the same solution; a rising diversity
  with stable reward signals broader coverage.
- Pairing Top-K metrics with ELBO/EUBO inside a `MultiMetricsModule` lets you
  track normalisation accuracy alongside sampling quality.
- Choose a fast distance function since the module computes pairwise
  distances across the top set each time you process new trajectories.

## What you get

| Metric key | Meaning |
| --- | --- |
| `top_{k}_reward` | Mean reward of the best `k` samples (rewards are converted to linear scale by exponentiating the logged rewards). |
| `top_{k}_diversity` | Average pairwise distance between the top `k` terminal states, measured with `distance_fn`. |

## Configuration

- `num_traj`: number of trajectories to sample each time you call `process`
  (should be ≥ max(top_k));
- `batch_size`: how many trajectories to roll out in parallel
- `top_k`: list of integers (e.g. `[10, 50, 100]`);
- `distance_fn`: callable that measures diversity between two terminal states.

## Quick start

> **Environment requirement:** supply a forward policy and distance function so batches of terminal states can be sampled and compared.

```python
import jax
import jax.numpy as jnp
import gfnx
from gfnx.utils.distances import hamming_distance

env = gfnx.HypergridEnvironment(reward_module=gfnx.EasyHypergridRewardModule())
params = env.init(jax.random.PRNGKey(0))

policy_params = {
    "forward_num_actions": env.action_space.n,
    "backward_num_actions": env.backward_action_space.n,
}


def uniform_forward_policy(rng_key, obs, policy_params):
    batch = obs.shape[0]
    forward_logits = jnp.zeros((batch, policy_params["forward_num_actions"]), dtype=jnp.float32)
    backward_logits = jnp.zeros((batch, policy_params["backward_num_actions"]), dtype=jnp.float32)
    info = {"forward_logits": forward_logits, "backward_logits": backward_logits}
    return forward_logits, info


def grid_hamming(lhs_state, rhs_state):
    return hamming_distance(lhs_state.state, rhs_state.state)


metrics = gfnx.metrics.TopKMetricsModule(
    env=env,
    fwd_policy_fn=uniform_forward_policy,
    num_traj=4096,
    batch_size=256,
    top_k=[10, 50, 200],
    distance_fn=grid_hamming,
)
state = metrics.init(jax.random.PRNGKey(1), metrics.InitArgs())

state = metrics.process(
    state,
    jax.random.PRNGKey(2),
    metrics.ProcessArgs(policy_params=policy_params, env_params=params),
)
report = metrics.get(state)
print(report["top_50_reward"], report["top_50_diversity"])
```
