        self.continuous = continuous

        self.norm_weights_()

        # whether to sample from the output discrete logits

        self.sample = sample

        # for genes -> expression network (the analogy is growing on me)

        self.accepts_latent = accepts_latent
        if accepts_latent:
            assert exists(dim_latent)

            self.encode_latent = nn.Linear(dim_latent, hidden_dim)
            self.encode_latent = maybe_weight_norm(self.encode_latent, name = 'weight', dim = None)
            self.post_norm_latent_added = nn.RMSNorm(hidden_dim)

        self.register_buffer('init_hiddens', torch.zeros(hidden_dim))

    def norm_weights_(self):
        if not self.weight_norm_linears:
            return

        for param in self.parameters():
            if not isinstance(param, nn.Linear):
                continue

            param.parametrization.weight.original.copy_(param.weight)

    def forward(
        self,
        x,
        hiddens = None,
        latent = None,
        sample_temperature = 1.
    ):
        assert xnor(exists(latent), self.accepts_latent)

        x = self.proj_in(x)
        x, forget = x[:-1], x[-1]

        x = F.silu(x)

        if exists(hiddens):
            past_mem = self.mem_norm(hiddens) * forget.sigmoid()
            x = x + past_mem

        if self.accepts_latent:


        self.num_selected = num_selected
        self.num_children = num_genes_per_island - num_selected
        self.tournament_size = tournament_size

        self.dim_gene = dim
        self.genes = nn.Parameter(l2norm(torch.randn(num_genes, dim)))

        self.split_islands = Rearrange('(i g) ... -> i g ...', i = num_islands)
        self.merge_islands = Rearrange('i g ... -> (i g) ...')

        self.num_elites = num_elites # todo - redo with affinity maturation algorithm from artificial immune system field
        self.mutation_std_dev = mutation_std_dev

        assert 0. <= num_frac_migrate <= 1.

        self.num_frac_migrate = num_frac_migrate
        self.migrate_genes_every = migrate_genes_every

        self.register_buffer('step', tensor(0))

    def __getitem__(self, idx):
        return l2norm(self.genes[idx])

    @torch.inference_mode()
    def evolve(
        self,
        fitnesses,
        temperature = 1.5
    ):
        device, num_selected = fitnesses.device, self.num_selected
        assert fitnesses.ndim == 1 and fitnesses.shape[0] == self.num_genes

        # split out the islands

        genes = self.genes
        num_islands = self.num_islands
        has_elites = self.num_elites > 0

        fitnesses = self.split_islands(fitnesses)
        genes = self.split_islands(genes)

        # local competition within each island


evolve

        self.dim_gene = dim
        self.genes = nn.Parameter(l2norm(torch.randn(num_genes, dim)))

        self.split_islands = Rearrange('(i g) ... -> i g ...', i = num_islands)
        self.merge_islands = Rearrange('i g ... -> (i g) ...')

        self.num_elites = num_elites # todo - redo with affinity maturation algorithm from artificial immune system field
        self.mutation_std_dev = mutation_std_dev

        assert 0. <= num_frac_migrate <= 1.

        self.num_frac_migrate = num_frac_migrate
        self.migrate_genes_every = migrate_genes_every

        self.register_buffer('step', tensor(0))

    def __getitem__(self, idx):
        return l2norm(self.genes[idx])

    @torch.inference_mode()
    def evolve(
        self,
        fitnesses,
        temperature = 1.5
    ):
        device, num_selected = fitnesses.device, self.num_selected
        assert fitnesses.ndim == 1 and fitnesses.shape[0] == self.num_genes

        # split out the islands

        genes = self.genes
        num_islands = self.num_islands
        has_elites = self.num_elites > 0

        fitnesses = self.split_islands(fitnesses)
        genes = self.split_islands(genes)

        # local competition within each island

        sorted_fitness, sorted_gene_ids = fitnesses.sort(dim = -1, descending = True)

        selected_gene_ids = sorted_gene_ids[:, :num_selected]
        selected_fitness = sorted_fitness[:, :num_selected]

        selected_gene_ids_for_gather = repeat(selected_gene_ids, '... -> ... d', d = self.dim_gene)

        selected_genes = genes.gather(1, selected_gene_ids_for_gather)

        # tournament

        num_children = self.num_children

        batch_randperm = torch.randn((num_islands, num_children, num_selected), device = device).argsort(dim = -1)
        tourn_ids = batch_randperm[..., :self.tournament_size]

        sorted_fitness = repeat(sorted_fitness, '... -> ... d', d = tourn_ids.shape[-1])

        tourn_fitness_ids = sorted_fitness.gather(1, tourn_ids)

        parent_ids = tourn_fitness_ids.topk(2, dim = -1).indices

        parent_ids = rearrange(parent_ids, 'i g parents -> i (g parents)')

        parent_ids = repeat(parent_ids, '... -> ... d', d = self.dim_gene)

        parents = selected_genes.gather(1, parent_ids)
        parents = rearrange(parents, 'i (g parents) d -> parents i g d', parents = 2)

        # cross over

        parent1, parent2 = parents

        children = parent1.lerp(parent2, (torch.randn_like(parent1) / temperature).sigmoid())

        # maybe migration

        if (
            divisible_by(self.step.item() + 1, self.migrate_genes_every) and
            self.num_islands > 1 and
            self.num_frac_migrate > 0.
        ):

            if has_elites:
                elites, selected_genes = selected_genes[:, :1], selected_genes[:, 1:]

            num_can_migrate = selected_genes.shape[1]

            num_migrate = max(1, num_can_migrate * self.num_frac_migrate)

            # fixed migration pattern - what i observe to work best, for now
            # todo - option to make it randomly selected with a mask

            selected_genes, migrants = selected_genes[:, -num_migrate:], selected_genes[:, :-num_migrate]

            migrants = torch.roll(migrants, 1, dims = (1,))

            selected_genes = cat((selected_genes, migrants), dim = 1)

            if has_elites:
                selected_genes = cat((elites, selected_genes), dim = 1)

        # concat children

        genes = torch.cat((selected_genes, children), dim = 1)

        # mutate

        if self.mutation_std_dev > 0:

            if has_elites:
                elites, genes = genes[:, :1], genes[:, 1:]

            genes.add_(torch.randn_like(genes) * self.mutation_std_dev)

            if has_elites:
                genes = torch.cat((elites, genes), dim = 1)

        genes = self.merge_islands(genes)

        self.genes.copy_(l2norm(genes))

        self.step.add_(1)

        return selected_gene_ids # return the selected gene ids, for the outer learning orchestrator to determine which mutations to accept

# main class

class BlackboxGradientSensing(Module):

    def __init__(
        self,
        actor: Module,
        *,
        accelerator: Accelerator | None = None,
        state_norm: StateNorm | Module | dict | None  = None,
        actor_is_recurrent = False,
        latent_gene_pool: LatentGenePool | dict | None = None,
        concat_latent_to_state = False, # if False, will pass in the latents as a kwarg `latent`, else try to concat it to the state
        crossover_every_step = 1,
        crossover_after_step = 0,
        num_env_interactions = 1000,
        noise_pop_size = 40,
        noise_std_dev: dict[str, float] | float = 0.1, # Appendix F in paper, appears to be constant for sim and real
        mutate_latent_genes = False,


        # for each reward and its anti, the number of standard deviations below the baseline they can be for acceptance

        self.num_std_below_mean_thres_accept = num_std_below_mean_thres_accept

        # the fraction of genes that must be above the given reward threshold as defined by the variable above, in order for said noise to be accepted

        assert 0 <= frac_genes_pass_thres_accept <= 1.
        self.frac_genes_pass_thres_accept = frac_genes_pass_thres_accept

        # expose a few computed variables

        self.num_episodes_per_learning_cycle = self.rollouts_for_machine.shape[0] * num_rollout_repeats * 2

        self.is_main = rank == 0

        # keep track of number of steps

        self.register_buffer('step', tensor(0))

    def sync_seed_(self):
        acc = self.accelerator
        rand_int = torch.randint(0, int(1e7), (), device = acc.device)
        seed = acc.reduce(rand_int)
        torch.manual_seed(seed.item())

    def log(self, **data):
        return self.accelerator.log(data, step = self.step.item())

    def save(self, path, overwrite = False):

        acc = self.accelerator

        acc.wait_for_everyone()

        if not acc.is_main_process:
            return

        path = Path(path)
        assert overwrite or not path.exists()

        pkg = dict(
            actor = self.actor.state_dict(),
            ema_actor = self.ema_actor.state_dict() if self.use_ema else None,
            state_norm = self.state_norm.state_dict() if self.use_state_norm else None,


save

        # expose a few computed variables

        self.num_episodes_per_learning_cycle = self.rollouts_for_machine.shape[0] * num_rollout_repeats * 2

        self.is_main = rank == 0

        # keep track of number of steps

        self.register_buffer('step', tensor(0))

    def sync_seed_(self):
        acc = self.accelerator
        rand_int = torch.randint(0, int(1e7), (), device = acc.device)
        seed = acc.reduce(rand_int)
        torch.manual_seed(seed.item())

    def log(self, **data):
        return self.accelerator.log(data, step = self.step.item())

    def save(self, path, overwrite = False):

        acc = self.accelerator

        acc.wait_for_everyone()

        if not acc.is_main_process:
            return

        path = Path(path)
        assert overwrite or not path.exists()

        pkg = dict(
            actor = self.actor.state_dict(),
            ema_actor = self.ema_actor.state_dict() if self.use_ema else None,
            state_norm = self.state_norm.state_dict() if self.use_state_norm else None,
            latents = self.gene_pool.state_dict() if exists(self.gene_pool) else None,
            step = self.step
        )

        torch.save(pkg, str(path))

    def load(self, path):
        path = Path(path)

        assert path.exists()

        pkg = torch.load(str(path), weights_only = True)

        self.actor.load_state_dict(pkg['actor'])

        if exists(self.gene_pool):
            self.gene_pool.load_state_dict(pkg['latents'])

        if self.use_ema:
            assert 'ema_actor' in pkg
            self.ema_actor.load_state_dict(pkg['ema_actor'])

        if self.use_state_norm:
            assert 'state_norm' in pkg
            self.state_norm.load_state_dict(pkg['state_norm'])


load

        acc = self.accelerator

        acc.wait_for_everyone()

        if not acc.is_main_process:
            return

        path = Path(path)
        assert overwrite or not path.exists()

        pkg = dict(
            actor = self.actor.state_dict(),
            ema_actor = self.ema_actor.state_dict() if self.use_ema else None,
            state_norm = self.state_norm.state_dict() if self.use_state_norm else None,
            latents = self.gene_pool.state_dict() if exists(self.gene_pool) else None,
            step = self.step
        )

        torch.save(pkg, str(path))

    def load(self, path):
        path = Path(path)

        assert path.exists()

        pkg = torch.load(str(path), weights_only = True)

        self.actor.load_state_dict(pkg['actor'])

        if exists(self.gene_pool):
            self.gene_pool.load_state_dict(pkg['latents'])

        if self.use_ema:
            assert 'ema_actor' in pkg
            self.ema_actor.load_state_dict(pkg['ema_actor'])

        if self.use_state_norm:
            assert 'state_norm' in pkg
            self.state_norm.load_state_dict(pkg['state_norm'])

        self.step.copy_(pkg['step'])

    def return_wrapped_actor(self) -> ActorWrapper | Module:

        if not (self.use_state_norm or self.actor_accepts_latents):
            return self.actor

        wrapped_actor = ActorWrapper(
            self.actor,
            state_norm = self.state_norm if self.use_state_norm else None,
            latent_gene_pool = self.gene_pool if self.actor_accepts_latents else None,
        )

        return wrapped_actor

    @torch.inference_mode()
    def forward(
        self,
        maybe_envs,
        num_env_interactions = None,
        show_progress = None,


        assert path.exists()

        pkg = torch.load(str(path), weights_only = True)

        self.actor.load_state_dict(pkg['actor'])

        if exists(self.gene_pool):
            self.gene_pool.load_state_dict(pkg['latents'])

        if self.use_ema:
            assert 'ema_actor' in pkg
            self.ema_actor.load_state_dict(pkg['ema_actor'])

        if self.use_state_norm:
            assert 'state_norm' in pkg
            self.state_norm.load_state_dict(pkg['state_norm'])

        self.step.copy_(pkg['step'])

    def return_wrapped_actor(self) -> ActorWrapper | Module:

        if not (self.use_state_norm or self.actor_accepts_latents):
            return self.actor

        wrapped_actor = ActorWrapper(
            self.actor,
            state_norm = self.state_norm if self.use_state_norm else None,
            latent_gene_pool = self.gene_pool if self.actor_accepts_latents else None,
        )

        return wrapped_actor

    @torch.inference_mode()
    def forward(
        self,
        maybe_envs,
        num_env_interactions = None,
        show_progress = None,
        seed = None,
        max_timesteps_per_interaction = None,
    ):
        max_timesteps_per_interaction = default(max_timesteps_per_interaction, self.max_timesteps)
        show_progress = default(show_progress, self.show_progress)
        num_env_interactions = default(num_env_interactions, self.num_env_interactions)

        (
            learning_rate,
            num_selected,
            noise_pop_size,
            num_rollout_repeats,
            factorized_noise,


        self.max_timesteps = max_timesteps

        # gene pool, another axis for scaling and bitter lesson

        num_genes = 1
        gene_pool = None

        if isinstance(latent_gene_pool, dict):
            gene_pool = LatentGenePool(**latent_gene_pool)
            gene_pool.to(device)

            num_genes = gene_pool.num_genes

        self.actor_accepts_latents = exists(gene_pool)
        self.concat_latent_to_state = concat_latent_to_state

        self.gene_pool = gene_pool
        self.num_genes = num_genes

        def default_calc_fitness(reward_stats):
            return reduce(reward_stats[:, 0], 'g s e -> g', 'mean')

        self.calc_fitness = default(calc_fitness, default_calc_fitness)

        self.crossover_every_step = crossover_every_step
        self.crossover_after_step = crossover_after_step

        # whether to do heritable mutations to the latent genes

        self.mutate_latent_genes = mutate_latent_genes

        self.latent_gene_noise_std_dev = latent_gene_noise_std_dev

        # optim

        optim_params = [named_params[param_name] for param_name in self.param_names]

        self.optim = optim_klass(optim_params, lr = learning_rate, betas = betas, **optim_kwargs)

        self.learning_rate = learning_rate
        self.weight_decay = weight_decay


hook

        self.crossover_after_step = crossover_after_step

        # whether to do heritable mutations to the latent genes

        self.mutate_latent_genes = mutate_latent_genes

        self.latent_gene_noise_std_dev = latent_gene_noise_std_dev

        # optim

        optim_params = [named_params[param_name] for param_name in self.param_names]

        self.optim = optim_klass(optim_params, lr = learning_rate, betas = betas, **optim_kwargs)

        self.learning_rate = learning_rate
        self.weight_decay = weight_decay

        # hooks

        if exists(optim_step_post_hook):
            def hook(*_):
                optim_step_post_hook()

            self.optim.register_step_post_hook(hook)

        # maybe state norm

        if isinstance(state_norm, dict):
            state_norm = StateNorm(**state_norm)

        self.use_state_norm = exists(state_norm)

        if self.use_state_norm:
            self.state_norm = state_norm
            state_norm.to(device)

        # progress bar

        self.show_progress = show_progress

        # number of interactions with environment for learning


            # mask out any noise candidates whose max reward mean is greater than baseline

            reward_threshold_accept = baseline_mean - reward_std * self.num_std_below_mean_thres_accept

            max_reward_mean = torch.amax(reward_mean, dim = -1)

            accept_mask = einx.greater_equal('g n, g -> g n', max_reward_mean, reward_threshold_accept)
            accept_mask = reduce(accept_mask.float(), 'g n -> n', 'mean') >= self.frac_genes_pass_thres_accept

            reward_deltas = einx.multiply('g n, n', reward_deltas, accept_mask.float()) # just zero out the reward deltas that do not pass the threshold

            # progress bar

            pbar_descriptions = [
                f'rewards: {baseline_mean.mean().item():.2f}',
                f'best: {reward_mean.amax().item():.2f}',
                f'accepted: {accept_mask.sum().item()} / {noise_pop_size}'
            ]

            def calculate_weights(reward_deltas, log = True):

                # get the top performing noise indices

                k = min(num_selected, reward_deltas.numel() // 2)

                ranked_reward_deltas, ranked_reward_indices = reward_deltas.abs().topk(k, dim = -1)

                # get the weights for the weighted sum of the topk noise according to eq (3)

                weights = ranked_reward_deltas / reward_std.clamp(min = 1e-3)

                # multiply by sign

                weights *= torch.sign(reward_deltas.gather(-1, ranked_reward_indices))

                if log:
                    pbar_descriptions.append(f'best delta: {ranked_reward_deltas.amax().item():.2f}')

                return weights, ranked_reward_indices

            # get the weights for update

            (
                weights,
                ranked_reward_indices,
            ) = calculate_weights(reduce(reward_deltas, 'g n -> n', 'mean'))

            # update the param one by one

            for name, noise in noises.items():

                param = params[name]

                # add the best "elite" noise directions weighted by eq (3)

                best_noises = noise[1:][ranked_reward_indices]

                update = einsum(best_noises, weights, 'n ..., n -> ...')


from __future__ import annotations

import pytest

import torch
from torch import nn

from blackbox_gradient_sensing.bgs import (
    BlackboxGradientSensing,
    Actor,
    LatentGenePool
)

# mock env

import numpy as np

class Sim:
    def reset(self, seed = None):
        return np.random.randn(5) # state

    def step(self, actions):
        return np.random.randn(5), np.random.randn(1), False # state, reward, done

# test BGS

@pytest.mark.parametrize('factorized_noise', (True, False))
@pytest.mark.parametrize('use_custom_actor', (True, False))
@pytest.mark.parametrize('use_state_norm', (True, False))
@pytest.mark.parametrize('actor_is_recurrent', (True, False))
@pytest.mark.parametrize('use_genetic_algorithm', (True, False))
@pytest.mark.parametrize('num_islands', (1, 2))
@pytest.mark.parametrize('mutate_latent_genes', (True, False))
@pytest.mark.parametrize('optimize_partial_network', (True, False))
@pytest.mark.parametrize('use_ema', (True, False))
@pytest.mark.parametrize('continuous', (True, False))
def test_bgs(
    factorized_noise,
    use_custom_actor,
    use_state_norm,
    actor_is_recurrent,
    use_genetic_algorithm,
    num_islands,
    mutate_latent_genes,
    optimize_partial_network,
    use_ema,
    continuous
):

    sim = Sim()

    actor = Actor(
        dim_state = 5,
        num_actions = 2,
        dim_latent = 32,
        continuous = continuous,
        accepts_latent = use_genetic_algorithm
    ) # actor with weight norm

    # test custom actor

    if use_custom_actor:
        actor = nn.Linear(5, 2)

        if (
            actor_is_recurrent or
            use_genetic_algorithm or
            optimize_partial_network or
            continuous
        ):
            pytest.skip()

    # maybe state norm

    state_norm = None

    if use_state_norm:
        state_norm = dict(dim_state = 5)

    # maybe genetic algorithm

    latent_gene_pool = None

    if use_genetic_algorithm:
        latent_gene_pool = dict(
            dim = 32,
            num_genes_per_island = 3,
            num_islands = num_islands,
            migrate_genes_every = 1,
            num_selected = 2,
            tournament_size = 2
        )

    # main evo strat orchestrator

    bgs = BlackboxGradientSensing(
        actor = actor,
        noise_pop_size = 2,      # number of noise perturbations
        num_selected = 1,         # topk noise selected for update
        num_rollout_repeats = 1,   # how many times to redo environment rollout, per noise
        torch_compile_actor = False,
        max_timesteps = 1,
        mutate_latent_genes = mutate_latent_genes,
        factorized_noise = factorized_noise,
        state_norm = state_norm,
        actor_is_recurrent = actor_is_recurrent,
        latent_gene_pool = latent_gene_pool,
        modules_to_optimize = {'to_embed'} if optimize_partial_network else None,
        cpu = True,
        use_ema = use_ema
    )

    bgs(sim, 2) # pass the simulation environment in - say for 100 interactions with env

    # after much training, save your learned policy (and optional state normalization) for finetuning on real env

    bgs.save('./actor-and-state-norm.pt', overwrite = True)

    bgs.load('./actor-and-state-norm.pt')

# test genetic algorithm

@pytest.mark.parametrize('num_islands', (1, 4))
def test_cross_over(
    num_islands
):
    gene_pool = LatentGenePool(
        dim = 32,
        num_genes_per_island = 3,
        num_islands = num_islands,
        migrate_genes_every = 1,
        num_selected = 2,
        tournament_size = 2
    )

    fitness = torch.randn(3 * num_islands)

    gene_pool.evolve(fitness)


Sim

from __future__ import annotations

import pytest

import torch
from torch import nn

from blackbox_gradient_sensing.bgs import (
    BlackboxGradientSensing,
    Actor,
    LatentGenePool
)

# mock env

import numpy as np

class Sim:
    def reset(self, seed = None):
        return np.random.randn(5) # state

    def step(self, actions):
        return np.random.randn(5), np.random.randn(1), False # state, reward, done

# test BGS

@pytest.mark.parametrize('factorized_noise', (True, False))
@pytest.mark.parametrize('use_custom_actor', (True, False))
@pytest.mark.parametrize('use_state_norm', (True, False))
@pytest.mark.parametrize('actor_is_recurrent', (True, False))
@pytest.mark.parametrize('use_genetic_algorithm', (True, False))
@pytest.mark.parametrize('num_islands', (1, 2))
@pytest.mark.parametrize('mutate_latent_genes', (True, False))
@pytest.mark.parametrize('optimize_partial_network', (True, False))
@pytest.mark.parametrize('use_ema', (True, False))
@pytest.mark.parametrize('continuous', (True, False))
def test_bgs(
    factorized_noise,
    use_custom_actor,
    use_state_norm,
    actor_is_recurrent,
    use_genetic_algorithm,
    num_islands,


class Sim:
    def reset(self, seed = None):
        return np.random.randn(5) # state

    def step(self, actions):
        return np.random.randn(5), np.random.randn(1), False # state, reward, done

# test BGS

@pytest.mark.parametrize('factorized_noise', (True, False))
@pytest.mark.parametrize('use_custom_actor', (True, False))
@pytest.mark.parametrize('use_state_norm', (True, False))
@pytest.mark.parametrize('actor_is_recurrent', (True, False))
@pytest.mark.parametrize('use_genetic_algorithm', (True, False))
@pytest.mark.parametrize('num_islands', (1, 2))
@pytest.mark.parametrize('mutate_latent_genes', (True, False))
@pytest.mark.parametrize('optimize_partial_network', (True, False))
@pytest.mark.parametrize('use_ema', (True, False))
@pytest.mark.parametrize('continuous', (True, False))
def test_bgs(
    factorized_noise,
    use_custom_actor,
    use_state_norm,
    actor_is_recurrent,
    use_genetic_algorithm,
    num_islands,
    mutate_latent_genes,
    optimize_partial_network,
    use_ema,
    continuous
):

    sim = Sim()

    actor = Actor(
        dim_state = 5,
        num_actions = 2,
        dim_latent = 32,
        continuous = continuous,
        accepts_latent = use_genetic_algorithm
    ) # actor with weight norm

    # test custom actor

    if use_custom_actor:
        actor = nn.Linear(5, 2)

        if (
            actor_is_recurrent or
            use_genetic_algorithm or
            optimize_partial_network or
            continuous
        ):
            pytest.skip()

    # maybe state norm

    state_norm = None

    if use_state_norm:
        state_norm = dict(dim_state = 5)

    # maybe genetic algorithm

    latent_gene_pool = None

    if use_genetic_algorithm:
        latent_gene_pool = dict(
            dim = 32,
            num_genes_per_island = 3,
            num_islands = num_islands,
            migrate_genes_every = 1,
            num_selected = 2,
            tournament_size = 2
        )

    # main evo strat orchestrator

    bgs = BlackboxGradientSensing(
        actor = actor,
        noise_pop_size = 2,      # number of noise perturbations
        num_selected = 1,         # topk noise selected for update
        num_rollout_repeats = 1,   # how many times to redo environment rollout, per noise
        torch_compile_actor = False,
        max_timesteps = 1,
        mutate_latent_genes = mutate_latent_genes,
        factorized_noise = factorized_noise,
        state_norm = state_norm,
        actor_is_recurrent = actor_is_recurrent,
        latent_gene_pool = latent_gene_pool,
        modules_to_optimize = {'to_embed'} if optimize_partial_network else None,
        cpu = True,
        use_ema = use_ema
    )

    bgs(sim, 2) # pass the simulation environment in - say for 100 interactions with env

    # after much training, save your learned policy (and optional state normalization) for finetuning on real env

    bgs.save('./actor-and-state-norm.pt', overwrite = True)

    bgs.load('./actor-and-state-norm.pt')

# test genetic algorithm

@pytest.mark.parametrize('num_islands', (1, 4))
def test_cross_over(
    num_islands
):
    gene_pool = LatentGenePool(
        dim = 32,
        num_genes_per_island = 3,
        num_islands = num_islands,
        migrate_genes_every = 1,
        num_selected = 2,
        tournament_size = 2
    )

    fitness = torch.randn(3 * num_islands)

    gene_pool.evolve(fitness)


        factorized_noise = factorized_noise,
        state_norm = state_norm,
        actor_is_recurrent = actor_is_recurrent,
        latent_gene_pool = latent_gene_pool,
        modules_to_optimize = {'to_embed'} if optimize_partial_network else None,
        cpu = True,
        use_ema = use_ema
    )

    bgs(sim, 2) # pass the simulation environment in - say for 100 interactions with env

    # after much training, save your learned policy (and optional state normalization) for finetuning on real env

    bgs.save('./actor-and-state-norm.pt', overwrite = True)

    bgs.load('./actor-and-state-norm.pt')

# test genetic algorithm

@pytest.mark.parametrize('num_islands', (1, 4))
def test_cross_over(
    num_islands
):
    gene_pool = LatentGenePool(
        dim = 32,
        num_genes_per_island = 3,
        num_islands = num_islands,
        migrate_genes_every = 1,
        num_selected = 2,
        tournament_size = 2
    )

    fitness = torch.randn(3 * num_islands)

    gene_pool.evolve(fitness)


reset

from __future__ import annotations

import pytest

import torch
from torch import nn

from blackbox_gradient_sensing.bgs import (
    BlackboxGradientSensing,
    Actor,
    LatentGenePool
)

# mock env

import numpy as np

class Sim:
    def reset(self, seed = None):
        return np.random.randn(5) # state

    def step(self, actions):
        return np.random.randn(5), np.random.randn(1), False # state, reward, done

# test BGS

@pytest.mark.parametrize('factorized_noise', (True, False))
@pytest.mark.parametrize('use_custom_actor', (True, False))
@pytest.mark.parametrize('use_state_norm', (True, False))
@pytest.mark.parametrize('actor_is_recurrent', (True, False))
@pytest.mark.parametrize('use_genetic_algorithm', (True, False))
@pytest.mark.parametrize('num_islands', (1, 2))
@pytest.mark.parametrize('mutate_latent_genes', (True, False))
@pytest.mark.parametrize('optimize_partial_network', (True, False))
@pytest.mark.parametrize('use_ema', (True, False))
@pytest.mark.parametrize('continuous', (True, False))
def test_bgs(
    factorized_noise,
    use_custom_actor,
    use_state_norm,


step

import pytest

import torch
from torch import nn

from blackbox_gradient_sensing.bgs import (
    BlackboxGradientSensing,
    Actor,
    LatentGenePool
)

# mock env

import numpy as np

class Sim:
    def reset(self, seed = None):
        return np.random.randn(5) # state

    def step(self, actions):
        return np.random.randn(5), np.random.randn(1), False # state, reward, done

# test BGS

@pytest.mark.parametrize('factorized_noise', (True, False))
@pytest.mark.parametrize('use_custom_actor', (True, False))
@pytest.mark.parametrize('use_state_norm', (True, False))
@pytest.mark.parametrize('actor_is_recurrent', (True, False))
@pytest.mark.parametrize('use_genetic_algorithm', (True, False))
@pytest.mark.parametrize('num_islands', (1, 2))
@pytest.mark.parametrize('mutate_latent_genes', (True, False))
@pytest.mark.parametrize('optimize_partial_network', (True, False))
@pytest.mark.parametrize('use_ema', (True, False))
@pytest.mark.parametrize('continuous', (True, False))
def test_bgs(
    factorized_noise,
    use_custom_actor,
    use_state_norm,
    actor_is_recurrent,
    use_genetic_algorithm,
    num_islands,


train.py

from adam_atan2_pytorch import AdoptAtan2
from blackbox_gradient_sensing import BlackboxGradientSensing, Actor

# sim environment, example using gymansium

import gymnasium as gym

continuous = True

sim = gym.make(
    'LunarLander-v3',
    render_mode = 'rgb_array',
    continuous = continuous
)

dim_state = sim.observation_space.shape[0]

# hyperparams

num_noises = 100      # number of noise perturbations, from which top is chosen for a weighted update - in paper this was 200 for sim, 3 for real
num_selected = 15     # number of elite perturbations chosen


from __future__ import annotations

import random
from random import randrange, choice

from math import sqrt
from copy import deepcopy
from functools import partial
from pathlib import Path
from typing import Callable

import numpy as np

import torch
from torch import cat, stack, nn, tensor, Tensor
import torch.nn.functional as F
from torch.nn import Module, ModuleList, Parameter
from torch.optim import Adam
from torch.func import functional_call

import torch.distributed as dist


from blackbox_gradient_sensing.bgs import (
    BlackboxGradientSensing,
    Actor,
    StateNorm,
    LatentGenePool
)


from __future__ import annotations

import pytest

import torch
from torch import nn

from blackbox_gradient_sensing.bgs import (
    BlackboxGradientSensing,
    Actor,
    LatentGenePool
)

# mock env

import numpy as np

class Sim:
    def reset(self, seed = None):
        return np.random.randn(5) # state