Parameter Space Searches

Each search produces a list of configurations. They are the stages of a Sweep: pass several to nest them, or call .sweep(fn) on one to sweep it alone.

A parameter may also map to a {label: sub-space} dict rather than to plain values. The label is what gets recorded in the results, while the sub-space it expands into supplies the arguments of the actual call. This is how you sweep, say, algorithm-specific hyperparameters under a single algorithm parameter.

jaxnasium.eval.GridSearch dataclass

Sets up a grid search over the given parameter values. All combinations of the parameter values are created in its configs.

Pass as a stage to Sweep, alone or chained with other searches. For only performing this grid search, GridSearch(...).sweep(fn, ...) is a shorthand for Sweep(fn, self, ...). See Sweep for more details.

Arguments: params: A dict of param names mapping to a non-empty list of values, or to a {label: sub-space} dict of branches. See Sweep for what a branch expands into.

Examples:

sweep = GridSearch(
    {
        "env_name": ["CartPole-v1", "Acrobot-v1"],
        "gamma": [0.95, 0.99, 0.999],
    }
).sweep(train)
for run in sweep:
    result = run()
    print(result.arguments, result.result)

Combined with a random search over learning rates:

sweep = Sweep(
    train,
    GridSearch({"env_name": ["CartPole-v1", "Acrobot-v1"]}),
    RandomSearch(
        {"learning_rate": (1e-4, 1e-2, "log")},
        num_samples=32,
        seed=jax.random.PRNGKey(0),
    ),
)
for run in sweep:
    result = run()
    print(result.arguments, result.result)

Branching, so each environment brings its own arguments along:

GridSearch(
    {
        "env_name": {
            "CartPole-v1": {
                "env": jym.make("CartPole-v1"),
                "total_timesteps": 100_000,
            },
            "Acrobot-v1": {
                "env": jym.make("Acrobot-v1"),
                "total_timesteps": 500_000,
            },
        }
    }
)

params: dict[str, list | dict] instance-attribute

configs() -> list[dict[str, Any]]

Returns every combination of the parameter values. For a branch, the combination holds its label.

sweep(fn: Callable, *, print_cost_estimate: bool = False) -> Sweep

Create a Sweep object with this grid search

jaxnasium.eval.RandomSearch dataclass

num_samples independent draws from the given parameters.

Pass as a stage to Sweep, alone or chained with other searches. For only performing this random search, RandomSearch(...).sweep(fn, ...) is a shorthand for Sweep(fn, self, ...). See Sweep for more details.

Examples:

sweep = RandomSearch(
    {"learning_rate": (1e-4, 1e-2, "log"), "gamma": [0.98, 0.99]},
    num_samples=64,
    seed=jax.random.PRNGKey(1),
).sweep(train)
for run in sweep:
    result = run()
    print(result.arguments, result.result)

Combined with a grid search over environments (same learning rates in each env):

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),
    ),
)
for run in sweep:
    result = run()
    print(result.arguments, result.result)

configs() -> list[dict[str, Any]]

Draw num_samples configurations

sweep(fn: Callable, *, print_cost_estimate: bool = False) -> Sweep

Create a Sweep with this search alone.

jaxnasium.eval.SobolSearch dataclass

Quasi-random draws that cover the parameter space more evenly than RandomSearch.

Pass as a stage to Sweep, alone or chained with other searches. For only performing this Sobol search, SobolSearch(...).sweep(fn, ...) is a shorthand for Sweep(fn, self, ...). See Sweep for more details.

Examples:

sweep = SobolSearch(
    {"learning_rate": (1e-4, 1e-2, "log"), "gamma": (0.9, 1.0)},
    num_samples=64,
    seed=jax.random.PRNGKey(1),
).sweep(train)
for run in sweep:
    result = run()
    print(result.arguments, result.result)

Combined with a grid search over environments:

sweep = Sweep(
    train,
    GridSearch({"env_name": ["CartPole-v1", "Acrobot-v1"]}),
    SobolSearch(
        {"learning_rate": (1e-4, 1e-2, "log"), "gamma": (0.9, 1.0)},
        num_samples=64,
        seed=jax.random.PRNGKey(1),
    ),
    seed=jax.random.PRNGKey(0),
)
for run in sweep:
    result = run()
    print(result.arguments, result.result)

configs() -> list[dict[str, Any]]

num_samples configurations spread evenly over the parameter space.

sweep(fn: Callable, *, print_cost_estimate: bool = False) -> Sweep

Create a Sweep with this search alone.

jaxnasium.eval.OneAtATimeSearch dataclass

One-factor-at-a-time (OFAT) search over discrete parameter values.

Sets the first value of each parameter list as the baseline, then varies one parameter at a time to each of its other values while holding the rest at baseline. Linear in the number of values.

Pass as a stage to Sweep, alone or chained with other searches. For only performing this search, OneAtATimeSearch(...).sweep(fn, ...) is a shorthand for Sweep(fn, self, ...).

Arguments: params: A dict of param names mapping to a non-empty list of values, or to a {label: sub-space} dict of branches (see Sweep). The first entry of each is the baseline value.

Examples:

sweep = OneAtATimeSearch(
    {
        "gamma": [0.99, 0.95, 0.999],
        "lr": [0.1, 0.01],
        # One tag in the results, two keyword arguments in the call.
        "network": {
            "mlp": {"actor_kwargs": MLP_KWARGS, "critic_kwargs": MLP_KWARGS},
            "simba": {
                "actor_kwargs": SIMBA_KWARGS,
                "critic_kwargs": SIMBA_KWARGS,
            },
        },
    }
).sweep(train)
for run in sweep:
    result = run()
    print(result.arguments, result.result)

configs() -> list[dict[str, Any]]

Returns the baseline config, then one config per non-baseline value of each parameter.

sweep(fn: Callable, *, print_cost_estimate: bool = False) -> Sweep

Create a Sweep with this search alone.