Sweeps

jaxnasium.eval.Sweep allows us to build a (parameter) sweep over any function by combining different search stages into a flat list of jobs.

The primary use case here is to create any kind of evaluation function which accepts some parameters, and then perform a hyperparameter sweep over said function. Below is a simple example, testing various learning rates in two environments.

import jax
import jaxnasium as jym
from jaxnasium.eval import GridSearch, RandomSearch, Sweep
from jaxnasium.algorithms import PPO

SEED = jax.random.key(0)

def simple_eval_function(env, learning_rate_start):
    env = jym.make(env)
    ppo = PPO(learning_rate_start=learning_rate_start)
    ppo, _metrics = ppo.train(SEED, env)
    return ppo.evaluate(SEED, env).mean()

sweep = Sweep(
    simple_eval_function,
    GridSearch({"env": ["CartPole-v1", "Acrobot-v1"]}),
    RandomSearch(
        {"learning_rate_start": (1e-4, 1e-2, "log")},
        num_samples=32,
        seed=SEED,
    ),
)

for run in sweep:
    result = run()
    print(result.arguments, result.result)

A standard algorithm evaluation function is provided in AlgorithmEvaluation. It accepts env, algorithm and any leftover keywords as hyperparameters.

Typically, one may want to launch jobs through a slurm system:

parser = argparse.ArgumentParser()
parser.add_argument(
    "index",
    type=int,
    nargs="?",
    help="sweep job to run; omit to print the number of jobs",
)
args = parser.parse_args()

sweep = Sweep(...)

if args.index is None:
    print(len(sweep))
    return

result = sweep[args.index]()
print(result.arguments)
print("mean return:", float(result.result["evaluation"].mean()))

python sweep.py          # prints number of jobs
#SBATCH --array=0-NUM_JOBS
python sweep.py $SLURM_ARRAY_TASK_ID

Full example

# sweep.py
import argparse
from pathlib import Path

import jax
from jaxnasium.algorithms import PPO, SAC
from jaxnasium.eval import AlgorithmEvaluation, GridSearch, SobolSearch, Sweep

OUTPUT_DIR = Path("results/sweep")
KEY = jax.random.key(0)
SEARCH_KEY = jax.random.PRNGKey(1)

DEFAULTS = {"total_timesteps": 1_000_000, "num_envs": 8, "log_function": None}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "index",
        type=int,
        nargs="?",
        help="sweep job to run; omit to print the number of jobs",
    )
    args = parser.parse_args()

    evaluation = AlgorithmEvaluation(
        jax.random.split(KEY, 10),
        batch_size=5,
        num_evaluations=50,
        return_train_metrics=True,
        save_path=OUTPUT_DIR,
    )

    sweep = Sweep(
        evaluation,
        GridSearch({"env": ["CartPole-v1", "Pendulum-v1"]}),
        GridSearch(
            {
                "algorithm": {
                    "PPO": [
                        {"algorithm": PPO(**DEFAULTS)},
                        SobolSearch(
                            {
                                "clip_coef": (0.1, 0.3),
                                "ent_coef_start": (1e-3, 1e-1, "log"),
                            },
                            num_samples=32,
                            seed=SEARCH_KEY,
                        ),
                    ],
                    "SAC": [
                        {"algorithm": SAC(**DEFAULTS)},
                        SobolSearch(
                            {
                                "tau": (1e-3, 5e-2, "log"),
                                "init_alpha": (0.05, 0.5),
                            },
                            num_samples=32,
                            seed=SEARCH_KEY,
                        ),
                    ],
                }
            }
        ),
        SobolSearch(
            {"learning_rate_start": (1e-4, 3e-3, "log"), "gamma": (0.95, 0.999)},
            num_samples=16,
            seed=SEARCH_KEY,
        ),
    )
    # 2 envs × (32 PPO + 32 SAC) × 16 shared draws = 2048 jobs.

    if args.index is None:
        print(len(sweep))
        return

    result = sweep[args.index]()
    print(result.arguments)
    print("mean return:", float(result.result["evaluation"].mean()))


if __name__ == "__main__":
    main()

python sweep.py          # prints 2048
#SBATCH --array=0-2047
python sweep.py $SLURM_ARRAY_TASK_ID

jaxnasium.eval.Sweep

Build a parameter sweep over function fn from one or more stages of parameter space search, like GridSearch, OneAtATimeSearch, RandomSearch, or SobolSearch.

For a single stage, RandomSearch(...).sweep(fn, ...) is equivalent.

Arguments: fn: The function to sweep. Must accept all swept params as keywords. If it implements with_labels(labels) -> fn, Sweep calls that before each job so branch labels can be recorded (as AlgorithmEvaluation does). *stages: GridSearch / OneAtATimeSearch / RandomSearch / SobolSearch stages to nest. print_cost_estimate: After creation, print the cost estimate of the first job to provide a rough estimate of the memory requirements of a job in the sweep. This triggers a compilation, so adds some overhead to the creation of the sweep.

Attributes: fn: Runs one SweepJob and returns its SweepResult. jobs: All jobs, to run in a loop or to dispatch by index (e.g. over a slurm job array).

Example:

sweep = Sweep(
    train,
    GridSearch({"env_name": ["CartPole-v1", "Acrobot-v1"]}),
    RandomSearch(
        {"learning_rate": (1e-4, 1e-2, "log")},
        num_samples=64,
        seed=jax.random.PRNGKey(1),  # the same 64 lrs in both envs
    ),
)
len(sweep)  # 2 envs * 64 learning rates
for run in sweep:
    result = run()
    print(result.arguments, result.result)
# Or dispatch one job: sweep[0]()

Branching over algorithms, each with its own hyperparameters:

sweep = Sweep(
    train,
    GridSearch(
        {
            "algorithm": {
                "PPO": [
                    {"algorithm": PPO()},
                    SobolSearch({"clip_coef": (0.1, 0.3)}, 32, seed=key),
                ],
                "SAC": [
                    {"algorithm": SAC()},
                    SobolSearch({"tau": (1e-3, 5e-2)}, 32, seed=key),
                ],
            }
        }
    ),
)
# Results record algorithm="PPO"; the run is called with algorithm=PPO().

Or in a slurm job array:

#SBATCH --array=0-5
python create_jobs_and_run_id.py --batch_idx $SLURM_ARRAY_TASK_ID

fn: Callable[[SweepJob], SweepResult] = run_job instance-attribute

jobs: list[SweepJob] = _expand_space(list(stages), [SweepJob({}, {})]) instance-attribute

__init__(fn: Callable, *stages: ParameterSpaceSearch, print_cost_estimate: bool = False)

__len__()

__iter__()

__repr__()

__getitem__(index: int | slice) -> partial[SweepResult] | Sweep

jaxnasium.eval.SweepJob dataclass

One dispatchable unit of work: a single call to the swept function.

Attributes: args: The keyword arguments the swept function is called with. labels: Branch labels recorded in place of what they expanded into. Overlay these on args to get the human-readable configuration.

args: dict[str, Any] instance-attribute

labels: dict[str, Any] instance-attribute

arguments: dict[str, Any] property

The input configuration as recorded

__init__(args: dict[str, Any], labels: dict[str, Any]) -> None

jaxnasium.eval.SweepResult dataclass

Simple container for the results of a sweep job. Most prominently, this contains the result and the arguments that produced it.

Attributes: start_time: Unix time at which the job containing this run started. end_time: Unix time at which the job containing this run finished. arguments: The parameter values that produced this result. Branch parameters hold their label, not the sub-space they expanded into. result: Whatever the swept function returned, for this configuration.

duration: float property

Wall-clock seconds spent on this job.

to_json(**kwargs: Any) -> str

Serialize to JSON, stringifying anything not natively serializable.

to_json_no_result(**kwargs: Any) -> str

Serialize to JSON, stringifying anything not natively serializable, excluding the result.

Algorithm Evaluation

AlgorithmEvaluation trains a Jaxnasium algorithm on an environment and evaluates the resulting agent, over one seed or many. It is an intended trial function for a Sweep, but may be used on its own.

import jax
from jaxnasium.eval import AlgorithmEvaluation

evaluation = AlgorithmEvaluation(
    seed=jax.random.split(jax.random.key(0), 10),
    env="CartPole-v1",
    algorithm="PPO",
    batch_size=5,
    save_path="results/",
)
returns = evaluation(learning_rate_start=3e-4)  # shape (10, num_evaluations)

Any keyword arguments given to the call that aren't seed, env or algorithm are treated as algorithm hyperparameters, which is what lets a sweep drive it directly.

jaxnasium.eval.AlgorithmEvaluation dataclass

Train and evaluate a jaxnasium algorithm on a given environment, for one seed or for many. A possible trial function for a Sweep.

Arguments: seed: Key for training and evaluation. Pass multiple keys to repeat the training and evaluation for each (e.g. AlgorithmEvaluation(jax.random.split(key, 10))). Optional at init: it can also be supplied when calling (e.g. from a Sweep), but not both. A seed is required by the time the evaluation actually runs. env: Environment to train on, or a name to pass to jym.make. eval_env: Environment to evaluate on, or a name to pass to jym.make. Defaults to None, in which case evaluation happens on env. algorithm: RLAlgorithm instance or name (e.g. "PPO"). batch_size: the batch size to jax.lax.map in case multiple seeds are given. num_evaluations: Episodes to evaluate the trained agent over. deterministic_policy: Whether to act greedily (instead of sampling) during evaluation. return_train_metrics: Also return the per-iteration training metrics. save_path: Where to write the results, or None (the default) to write nothing and leave saving to the caller. This will save to / and writes 4 files: - results.npy: the evaluation results, shape (num_seeds, num_evaluations) - train_curve.npy: the training metrics, shape (num_seeds, num_iterations, ...) - input_parameters.json: the input parameters to the evaluation - full_parameters.json: the full parameters used for the evaluation. save_as_zip: If True, while save_path is not None, will zip the results folder and write it to save_path.

Example:

key = jax.random.key(0)
evaluation = AlgorithmEvaluation(
    jax.random.split(key, 10), env="CartPole-v1", batch_size=5
)
sweep = Sweep(
    evaluation,
    GridSearch({"algorithm": ["PPO", "DQN"]}),
    RandomSearch(
        {"learning_rate": (1e-4, 1e-2, "log")},
        num_samples=64,
        seed=jax.random.PRNGKey(1),
    ),
)
# 2 algorithms * 64 learning rates = 128 jobs, each of 10 seeds, 5 at a time.

__call__(**kwargs: Any) -> Any

Train and evaluate, once per seed.

Returns a dict {"evaluation": evaluation_returns, "train_metrics": train_metrics} if return_train_metrics is True, otherwise just the evaluation returns. With several seeds, every leaf gains a leading axis over them.

context(**kwargs: Any) -> dict[str, Any]

JSON-ready metadata for a configuration, to pair with its result. Used to retrieve the full algorithm and env parameters that were used for this run.

with_labels(labels: dict[str, Any]) -> AlgorithmEvaluation

Attach recorded labels for the saved config that won't be passed as call arguments.