PopuLoRA

Quick start

import torch
import torch.nn as nn
from populora import Population

model = nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(8, 1))

pop = Population(model, pop_size = 16, low_rank = 4, lora_targets = ['0', '2'])

state = torch.randn(1, 4, 2)
preds = pop(state, all_individuals = True)  # one routed forward over every individual

labels = torch.randn(1, 4, 1)
fitnesses = -((preds - labels) ** 2).reshape(16, -1).mean(dim = -1)

result = pop.select('deterministic', fitnesses, survive_frac = 0.5)

parents = pop.select_parents(
    'tournament',
    fitnesses,
    num_children = len(result.selected_out_indices),
    culled = result.selected_out_indices  # parents come from the survivors
)

pop.crossover_('average', parents, result.selected_out_indices)  # offspring overwrite the culled
pop.mutate_('full_gaussian', individuals = result.selected_out_indices)

model = pop.merge_(fitnesses.argmax())  # merge the best individual back in

pop.save_individual('best.pt', fitnesses.argmax())  # or save just the best individual's weights

pop.evolve_(fitnesses) runs selection, parent selection, crossover, and mutation in one step. Batch evaluation also supports pop(x, individuals = [ids]) to route each sample to its own individual.

Environment evolution

evolve_with_env evolves a population against any MDP-style environment in one call — gymnasium, dm_control, isaac, maniskill, pybullet, pufferlib, or any simulator — and returns the merged best policy.

from torch import nn
from dm_control import suite
from populora import evolve_with_env

policy, history = evolve_with_env(
    [suite.load('cartpole', 'balance') for _ in range(16)],  # envs, a list, a vector env, or a factory
    nn.Sequential(nn.Linear(5, 32), nn.ReLU(), nn.Linear(32, 2)),
    pop_size = 16,
    low_rank = 16,
    action = lambda logits: logits.argmax(-1) * 2 - 1,
    num_generations = 25,
    horizon = 1000,
    seed = 0,
    progress = True,
    return_history = True  # per-generation best / mean
)

lora_targets auto-discovers every Linear layer when omitted. Pass target_fitness to stop early, checkpoint_dir to write latest.pt (every checkpoint_every generations) and best.pt (on new bests), and resume = True to pick up from the latest checkpoint.

Continuous policies emit actions on (-1, 1) — pass to_range = (-2., 2.) (or any env action range) to evolve_with_env / interact_with_env and the actions are rescaled before stepping the env. An action fn built with make_action carries its own range, e.g. a beta with beta_rescale_neg_one_one = False emits on (0, 1) instead — the interactor rescales from that automatically. interact_with_env(from_range = ...) is the fallback for custom action functions, and rescale_from_range_to_range is exported too.

For a custom loop, interact_with_env exposes the underlying EnvInteractor — one routed forward over the active slots per timestep, distributed across ranks under torchrun:

from torch import nn
from dm_control import suite
from populora import interact_with_env

interactor = interact_with_env([suite.load('cartpole', 'balance') for _ in range(16)])

backbone = nn.Sequential(nn.Linear(5, 32), nn.ReLU(), nn.Linear(32, 2))
population = interactor.population(backbone, pop_size = 16, low_rank = 16)

for gen in range(25):
    fitnesses = interactor.evaluate(
        population,
        action = lambda logits: logits.argmax(-1) * 2 - 1,
        horizon = 1000
    )

    population.evolve_(fitnesses)

policy = population.merge_(fitnesses.argmax())

Custom fitness functions may take (population, individuals) (batched), (population, idx) (per index), or (population) (all at once) — detected automatically.