x-evolution 0.1.30__tar.gz → 0.1.32__tar.gz

{x_evolution-0.1.30 → x_evolution-0.1.32}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.30
+Version: 0.1.32
 Summary: x-evolution
 Project-URL: Homepage, https://pypi.org/project/x-evolution/
 Project-URL: Repository, https://github.com/lucidrains/x-evolution
@@ -37,6 +37,7 @@ Requires-Python: >=3.9
 Requires-Dist: accelerate
 Requires-Dist: beartype
 Requires-Dist: einops>=0.8.0
+Requires-Dist: torch-einops-utils
 Requires-Dist: torch>=2.4
 Requires-Dist: x-mlps-pytorch>=0.2.0
 Requires-Dist: x-transformers>=2.11.23
@@ -52,6 +53,8 @@ Description-Content-Type: text/markdown
 
 Implementation of various evolutionary algorithms, starting with [evolutionary strategies](https://blog.otoro.net/2017/10/29/visual-evolution-strategies/)
 
+[explained by @bycloud](https://www.youtube.com/watch?v=lLkE9w1NJs0)
+
 ## Install
 
 ```bash
@@ -95,19 +98,29 @@ evo_strat()
 
 ## Distributed
 
-Using the CLI from 🤗 
+Using the CLI from 🤗
 
 ```shell
-$ accelerate config
+$ uv run accelerate config
 ```
 
 Then
 
 ```shell
-$ accelerate launch train.py
+$ uv run accelerate launch train.py
+```
+
+Or with `torchrun` (which can do multiple CPU processes for testing)
+
+```shell
+$ uv run torchrun --nproc_per_node=<N> train.py --cpu True
 ```
 
-For gymnasium simulations, first run `pip install '[examples]'`
+For a quick demo on LunarLander with 8 processes
+
+```shell
+$ uv run torchrun --nproc_per_node=8 train_lunar.py --cpu True --noise_population_size 100 --fitness_to_weighted_factor centered_rank
+```
 
 ## Citations
 
@@ -124,25 +137,25 @@ For gymnasium simulations, first run `pip install '[examples]'`
 
 ```bibtex
 @misc{sarkar2025evolutionstrategieshyperscale,
-    title   = {Evolution Strategies at the Hyperscale}, 
+    title   = {Evolution Strategies at the Hyperscale},
     author  = {Bidipta Sarkar and Mattie Fellows and Juan Agustin Duque and Alistair Letcher and Antonio León Villares and Anya Sims and Dylan Cope and Jarek Liesen and Lukas Seier and Theo Wolf and Uljad Berdica and Alexander David Goldie and Aaron Courville and Karin Sevegnani and Shimon Whiteson and Jakob Nicolaus Foerster},
     year    = {2025},
     eprint  = {2511.16652},
     archivePrefix = {arXiv},
     primaryClass = {cs.LG},
-    url     = {https://arxiv.org/abs/2511.16652}, 
+    url     = {https://arxiv.org/abs/2511.16652},
 }
 ```
 
 ```bibtex
 @misc{fortunato2019noisynetworksexploration,
-    title   = {Noisy Networks for Exploration}, 
+    title   = {Noisy Networks for Exploration},
     author  = {Meire Fortunato and Mohammad Gheshlaghi Azar and Bilal Piot and Jacob Menick and Ian Osband and Alex Graves and Vlad Mnih and Remi Munos and Demis Hassabis and Olivier Pietquin and Charles Blundell and Shane Legg},
     year    = {2019},
     eprint  = {1706.10295},
     archivePrefix = {arXiv},
     primaryClass = {cs.LG},
-    url     = {https://arxiv.org/abs/1706.10295}, 
+    url     = {https://arxiv.org/abs/1706.10295},
 }
 ```
 

{x_evolution-0.1.30 → x_evolution-0.1.32}/README.md

@@ -2,6 +2,8 @@
 
 Implementation of various evolutionary algorithms, starting with [evolutionary strategies](https://blog.otoro.net/2017/10/29/visual-evolution-strategies/)
 
+[explained by @bycloud](https://www.youtube.com/watch?v=lLkE9w1NJs0)
+
 ## Install
 
 ```bash
@@ -45,19 +47,29 @@ evo_strat()
 
 ## Distributed
 
-Using the CLI from 🤗 
+Using the CLI from 🤗
 
 ```shell
-$ accelerate config
+$ uv run accelerate config
 ```
 
 Then
 
 ```shell
-$ accelerate launch train.py
+$ uv run accelerate launch train.py
+```
+
+Or with `torchrun` (which can do multiple CPU processes for testing)
+
+```shell
+$ uv run torchrun --nproc_per_node=<N> train.py --cpu True
 ```
 
-For gymnasium simulations, first run `pip install '[examples]'`
+For a quick demo on LunarLander with 8 processes
+
+```shell
+$ uv run torchrun --nproc_per_node=8 train_lunar.py --cpu True --noise_population_size 100 --fitness_to_weighted_factor centered_rank
+```
 
 ## Citations
 
@@ -74,25 +86,25 @@ For gymnasium simulations, first run `pip install '[examples]'`
 
 ```bibtex
 @misc{sarkar2025evolutionstrategieshyperscale,
-    title   = {Evolution Strategies at the Hyperscale}, 
+    title   = {Evolution Strategies at the Hyperscale},
     author  = {Bidipta Sarkar and Mattie Fellows and Juan Agustin Duque and Alistair Letcher and Antonio León Villares and Anya Sims and Dylan Cope and Jarek Liesen and Lukas Seier and Theo Wolf and Uljad Berdica and Alexander David Goldie and Aaron Courville and Karin Sevegnani and Shimon Whiteson and Jakob Nicolaus Foerster},
     year    = {2025},
     eprint  = {2511.16652},
     archivePrefix = {arXiv},
     primaryClass = {cs.LG},
-    url     = {https://arxiv.org/abs/2511.16652}, 
+    url     = {https://arxiv.org/abs/2511.16652},
 }
 ```
 
 ```bibtex
 @misc{fortunato2019noisynetworksexploration,
-    title   = {Noisy Networks for Exploration}, 
+    title   = {Noisy Networks for Exploration},
     author  = {Meire Fortunato and Mohammad Gheshlaghi Azar and Bilal Piot and Jacob Menick and Ian Osband and Alex Graves and Vlad Mnih and Remi Munos and Demis Hassabis and Olivier Pietquin and Charles Blundell and Shane Legg},
     year    = {2019},
     eprint  = {1706.10295},
     archivePrefix = {arXiv},
     primaryClass = {cs.LG},
-    url     = {https://arxiv.org/abs/1706.10295}, 
+    url     = {https://arxiv.org/abs/1706.10295},
 }
 ```
 

{x_evolution-0.1.30 → x_evolution-0.1.32}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-evolution"
-version = "0.1.30"
+version = "0.1.32"
 description = "x-evolution"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -27,6 +27,7 @@ dependencies = [
     "beartype",
     "einops>=0.8.0",
     "torch>=2.4",
+    "torch-einops-utils",
     "x-mlps-pytorch>=0.2.0",
     "x-transformers>=2.11.23"
 ]
@@ -65,3 +66,11 @@ allow-direct-references = true
 
 [tool.hatch.build.targets.wheel]
 packages = ["x_evolution"]
+
+[tool.hatch.build.targets.sdist]
+include = [
+    "x_evolution",
+    "README.md",
+    "LICENSE",
+    "pyproject.toml",
+]

{x_evolution-0.1.30 → x_evolution-0.1.32}/x_evolution/x_evolution.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Callable
+from typing import Callable, Literal
 
 from math import ceil
 from pathlib import Path
@@ -27,6 +27,7 @@ from x_mlps_pytorch.noisable import (
 )
 
 from einops import pack
+from torch_einops_utils import pack_with_inverse
 
 # constants
 
@@ -46,15 +47,31 @@ def identity(t):
 def divisible_by(num, den):
     return (num % den) == 0
 
-def normalize(t, eps = 1e-6):
-    return F.layer_norm(t, t.shape, eps = eps)
-
 def accum_grad_(t, value):
     if not exists(t.grad):
         t.grad = value.clone()
     else:
         t.grad.add_(value)
 
+# fitness weighting factors
+
+def normalize(t, eps = 1e-6):
+    return F.layer_norm(t, t.shape, eps = eps)
+
+def centered_rank_normalize(fitnesses):
+    dtype, device, n = fitnesses.dtype, fitnesses.device, fitnesses.numel()
+    flat_fitnesses, inverse = pack_with_inverse(fitnesses, '*')
+    ranks = flat_fitnesses.argsort().argsort().to(dtype = dtype)
+    ranks = (ranks / (n - 1)) - 0.5
+    return inverse(ranks)
+
+FITNESS_WEIGHT_FACTORS = dict(
+    normalize = normalize,
+    centered_rank = centered_rank_normalize
+)
+
+FitnessWeightFactorType = Literal['normalize', 'centered_rank']
+
 # class
 
 class EvoStrategy(Module):
@@ -89,7 +106,7 @@ class EvoStrategy(Module):
         sigma_scheduler_klass: type[LRScheduler] | None = None,
         sigma_scheduler_kwargs: dict = dict(),
         transform_fitness: Callable = identity,
-        fitness_to_weighted_factor: Callable[[Tensor], Tensor] = normalize,
+        fitness_to_weighted_factor: Callable[[Tensor], Tensor] | FitnessWeightFactorType = 'normalize',
         checkpoint_every = None,            # saving every number of generations
         checkpoint_path = './checkpoints',
         cpu = False,
@@ -97,6 +114,7 @@ class EvoStrategy(Module):
         accelerator: Accelerator | None = None,
         accelerate_kwargs: dict = dict(),
         reject_generation_fitnesses_if: Callable[[Tensor], bool] | None = None,
+        on_result: Callable | None = None,
         vectorized = False,
         vector_size: int | None = None,
         sync_on_init = True
@@ -106,6 +124,7 @@ class EvoStrategy(Module):
 
         self.vectorized = vectorized
         self.vector_size = vector_size
+        self.on_result = on_result
 
         if not exists(accelerator):
             accelerator = Accelerator(cpu = cpu, **accelerate_kwargs)
@@ -179,7 +198,7 @@ class EvoStrategy(Module):
         # dtypes
 
         self.param_dtypes = {name: param.dtype for name, param in named_parameters_dict.items()}
- 
+
         # hyperparameters
 
         self.noise_population_size = noise_population_size
@@ -191,6 +210,11 @@ class EvoStrategy(Module):
 
         self.transform_fitness = transform_fitness # a function that gets called before converting to weights for the weighted noise update - eventually get rank normalization
 
+        if isinstance(fitness_to_weighted_factor, str):
+            if fitness_to_weighted_factor not in FITNESS_WEIGHT_FACTORS:
+                raise ValueError(f'fitness_to_weighted_factor must be one of {list(FITNESS_WEIGHT_FACTORS.keys())}')
+            fitness_to_weighted_factor = FITNESS_WEIGHT_FACTORS[fitness_to_weighted_factor]
+
         self.fitness_to_weighted_factor = fitness_to_weighted_factor
 
         self.mirror_sampling = mirror_sampling # mirror / antithetical sampling - reducing variance by doing positive + negative of noise and subtracting
@@ -522,7 +546,10 @@ class EvoStrategy(Module):
 
                 if not self.mirror_sampling:
                     fitnesses.append(fitness)
-                    pbar.set_postfix(reward = f'{fitness:.3f}')
+
+                    if exists(self.on_result):
+                        self.on_result(fitness, pbar)
+
                     continue
 
                 # handle mirror sampling
@@ -530,7 +557,10 @@ class EvoStrategy(Module):
                 fitness_mirrored = get_fitness(negate = True)
 
                 fitnesses.append([fitness, fitness_mirrored])
-                pbar.set_postfix(reward = f'{(fitness + fitness_mirrored) / 2:.3f}')
+                avg_fitness = (fitness + fitness_mirrored) / 2
+
+                if exists(self.on_result):
+                    self.on_result(avg_fitness, pbar)
 
             pbar.close()
 

x_evolution-0.1.30/.github/workflows/python-publish.yml

@@ -1,36 +0,0 @@
-# This workflow will upload a Python Package using Twine when a release is created
-# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
-
-# This workflow uses actions that are not certified by GitHub.
-# They are provided by a third-party and are governed by
-# separate terms of service, privacy policy, and support
-# documentation.
-
-name: Upload Python Package
-
-on:
-  release:
-    types: [published]
-
-jobs:
-  deploy:
-
-    runs-on: ubuntu-latest
-
-    steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python
-      uses: actions/setup-python@v2
-      with:
-        python-version: '3.x'
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install build
-    - name: Build package
-      run: python -m build
-    - name: Publish package
-      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
-      with:
-        user: __token__
-        password: ${{ secrets.PYPI_API_TOKEN }}

x_evolution-0.1.30/.github/workflows/test.yml

@@ -1,21 +0,0 @@
-name: Pytest
-on: [push, pull_request]
-
-jobs:
-  build:
-
-    runs-on: ubuntu-latest
-
-    steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python 3.10
-      uses: actions/setup-python@v5
-      with:
-        python-version: "3.10"
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        python -m pip install -e .[test]
-    - name: Test with pytest
-      run: |
-        python -m pytest tests/

x_evolution-0.1.30/.gitlab-ci.yml

@@ -1,27 +0,0 @@
-image: python:3.10
-
-stages:
-  - test
-  - publish
-
-test:
-  stage: test
-  before_script:
-    - python -m pip install --upgrade pip
-    - python -m pip install -e .[test]
-  script:
-    - python -m pytest tests/
-  only:
-    - pushes
-    - merge_requests
-
-pypi:
-  stage: publish
-  before_script:
-    - python -m pip install --upgrade pip
-    - pip install build twine
-  script:
-    - python -m build
-    - TWINE_PASSWORD=$PYPI_API_TOKEN TWINE_USERNAME=__token__ python -m twine upload dist/*
-  only:
-    - tags

x_evolution-0.1.30/scripts/install-humanoid.sh

@@ -1,2 +0,0 @@
-#!/bin/bash
-apt-get update && apt-get install -y libegl1 libgl1-mesa-glx libosmesa6-dev xvfb

x_evolution-0.1.30/tests/test_evolution.py

@@ -1,70 +0,0 @@
-import pytest
-param = pytest.mark.parametrize
-
-import torch
-from x_mlps_pytorch import MLP
-
-model = MLP(8, 16, 4)
-
-@param('params_to_optimize', (None, [model.layers[1].weight]))
-@param('use_optimizer', (False, True))
-@param('noise_low_rank', (None, 1))
-@param('mirror_sampling', (False, True))
-@param('multi_models', (False, True))
-@param('learned_sigma', (False, True))
-@param('use_sigma_optimizer', (False, True))
-def test_evo_strat(
-    params_to_optimize,
-    use_optimizer,
-    noise_low_rank,
-    mirror_sampling,
-    multi_models,
-    learned_sigma,
-    use_sigma_optimizer
-):
-    from random import randrange
-
-    from x_evolution.x_evolution import EvoStrategy
-
-    to_optim = model
-
-    if multi_models:
-        to_optim = [model, MLP(8, 1)]
-
-    evo_strat = EvoStrategy(
-        to_optim,
-        environment = lambda model: float(randrange(100)),
-        num_generations = 2,
-        learned_noise_scale = learned_sigma,
-        params_to_optimize = params_to_optimize,
-        use_optimizer = use_optimizer,
-        noise_low_rank = noise_low_rank,
-        mirror_sampling = mirror_sampling,
-        use_sigma_optimizer = use_sigma_optimizer
-    )
-
-    evo_strat('evolve', 1)
-    evo_strat('more.evolve', 1)
-
-    fitnesses = evo_strat('more.evolve', 2, rollback_model_at_end = True)
-
-@param('vector_size', (2, 4))
-def test_evo_strat_vectorized(vector_size):
-    from x_evolution.x_evolution import EvoStrategy
-
-    model = MLP(8, 16, 4)
-
-    def environment(model):
-        # mock a vectorized environment returning multiple fitness scores
-        return torch.randn(vector_size)
-
-    evo_strat = EvoStrategy(
-        model,
-        environment = environment,
-        num_generations = 2,
-        vectorized = True,
-        vector_size = vector_size,
-        noise_population_size = 4
-    )
-
-    evo_strat()

x_evolution-0.1.30/train_ant_maze.py

@@ -1,211 +0,0 @@
-# /// script
-# dependencies = [
-#     "fire",
-#     "gymnasium[mujoco]>=1.0.0",
-#     "gymnasium-robotics",
-#     "moviepy",
-#     "tqdm",
-#     "x-evolution>=0.1.27",
-#     "x-mlps-pytorch"
-# ]
-# ///
-
-import fire
-from shutil import rmtree
-import gymnasium as gym
-import gymnasium_robotics
-import numpy as np
-from tqdm import tqdm
-
-import torch
-from torch import nn
-from torch.nn import Module, GRU, Linear
-import torch.nn.functional as F
-
-# functions
-
-def exists(v):
-    return v is not None
-
-def softclamp(t, value):
-    return (t / value).tanh() * value
-
-class AntMazeEnvironment(Module):
-    def __init__(
-        self,
-        env_id = 'AntMaze_UMazeDense-v5',
-        video_folder = './recordings_ant_maze',
-        render_every_eps = 100,
-        max_steps = 1000,
-        repeats = 1,
-        verbose = False
-    ):
-        super().__init__()
-
-        self.verbose = verbose
-
-        env = gym.make(env_id, render_mode = 'rgb_array')
-
-        self.env = env
-        self.max_steps = max_steps
-        self.repeats = repeats
-        self.video_folder = video_folder
-        self.render_every_eps = render_every_eps
-
-    def pre_main_callback(self):
-        rmtree(self.video_folder, ignore_errors = True)
-
-        self.env = gym.wrappers.RecordVideo(
-            env = self.env,
-            video_folder = self.video_folder,
-            name_prefix = 'recording',
-            episode_trigger = lambda eps_num: (eps_num % self.render_every_eps) == 0,
-            disable_logger = True
-        )
-
-    def forward(self, model):
-
-        device = next(model.parameters()).device
-
-        seed = torch.randint(0, int(1e6), ())
-
-        num_envs = 1
-        cum_reward = torch.zeros(num_envs, device = device)
-
-        for _ in range(self.repeats):
-            obs, _ = self.env.reset(seed = seed.item())
-
-            step = 0
-            hiddens = None
-            
-            dones = torch.zeros(num_envs, device = device, dtype = torch.bool)
-
-            while step < self.max_steps and not dones.all():
-
-                state = obs['observation']
-                goal = obs['desired_goal']
-
-                state_torch = torch.from_numpy(state).float().to(device)
-                goal_torch = torch.from_numpy(goal).float().to(device)
-
-                state_goal = torch.cat((state_torch, goal_torch), dim = -1)
-
-                action_logits, hiddens = model(state_goal, hiddens)
-
-                mean, log_var = action_logits.chunk(2, dim = -1)
-
-                std = (0.5 * softclamp(log_var, 5.)).exp()
-                sampled = mean + torch.randn_like(mean) * std
-                action = sampled.tanh()
-
-                next_obs, reward, truncated, terminated, info = self.env.step(action.detach().cpu().numpy())
-
-                reward_np = np.array(reward) if not isinstance(reward, np.ndarray) else reward
-                total_reward_base = torch.from_numpy(reward_np).float().to(device)
-
-                exploration_bonus = std.mean(dim = -1) * 0.01
-                penalize_extreme_actions = (mean.abs() > 1.).float().mean(dim = -1) * 0.01
-
-                total_reward = total_reward_base + exploration_bonus - penalize_extreme_actions
-
-                mask = (~dones).float()
-                cum_reward += total_reward * mask
-
-                dones_np = np.array(truncated | terminated) if not isinstance(truncated | terminated, np.ndarray) else (truncated | terminated)
-                dones |= torch.from_numpy(dones_np).to(device)
-
-                step += 1
-                obs = next_obs
-
-        return cum_reward.item() / self.repeats
-
-# model
-
-from x_mlps_pytorch.residual_normed_mlp import ResidualNormedMLP
-
-class Model(Module):
-
-    def __init__(self, state_dim = 27 + 2, action_dim = 8):
-        super().__init__()
-
-        self.deep_mlp = ResidualNormedMLP(
-            dim_in = state_dim,
-            dim = 256,
-            depth = 8,
-            residual_every = 2
-        )
-
-        self.gru = GRU(256, 256, batch_first = True)
-
-        self.to_pred = Linear(256, action_dim * 2, bias = False)
-
-    def forward(self, state, hiddens = None):
-
-        x = self.deep_mlp(state)
-
-        is_batched = x.ndim > 1
-
-        if not is_batched:
-            x = x.unsqueeze(0)
-
-        x = x.unsqueeze(-2)
-        gru_out, hiddens = self.gru(x, hiddens)
-        x = x + gru_out
-        x = x.squeeze(-2)
-
-        if not is_batched:
-            x = x.squeeze(0)
-
-        return self.to_pred(x), hiddens
-
-# evo strategy
-
-from x_evolution import EvoStrategy
-from torch.optim.lr_scheduler import CosineAnnealingLR
-
-def main(
-    env_id = 'AntMaze_UMazeDense-v5',
-    num_generations = 50_000,
-    max_steps = 250,
-    noise_population_size = 150,
-    verbose = True
-):
-    # Determine dims from temp env
-    temp_env = gym.make(env_id)
-    obs_space = temp_env.observation_space
-    action_space = temp_env.action_space
-    
-    state_dim = obs_space['observation'].shape[0] + obs_space['desired_goal'].shape[0]
-    action_dim = action_space.shape[0]
-    temp_env.close()
-
-    model = Model(state_dim = state_dim, action_dim = action_dim)
-
-    evo_strat = EvoStrategy(
-        model,
-        environment = AntMazeEnvironment(
-            env_id = env_id,
-            max_steps = max_steps,
-            repeats = 1,
-            render_every_eps = 200,
-            verbose = verbose
-        ),
-        num_generations = num_generations,
-        noise_population_size = noise_population_size,
-        noise_low_rank = 1,
-        noise_scale = 1e-2,
-        noise_scale_clamp_range = (5e-3, 2e-2),
-        learned_noise_scale = True,
-        use_sigma_optimizer = True,
-        learning_rate = 3e-4,
-        noise_scale_learning_rate = 1e-4,
-        use_scheduler = True,
-        scheduler_klass = CosineAnnealingLR,
-        scheduler_kwargs = dict(T_max = num_generations),
-        verbose = verbose
-    )
-
-    evo_strat()
-
-if __name__ == '__main__':
-    fire.Fire(main)

x_evolution-0.1.30/train_humanoid.py

@@ -1,215 +0,0 @@
-# /// script
-# dependencies = [
-#     "fire",
-#     "gymnasium[mujoco]>=1.0.0",
-#     "gymnasium[other]",
-#     "x-evolution>=0.0.20",
-#     "x-mlps-pytorch"
-# ]
-# ///
-
-# import os
-# os.environ["NCCL_P2P_DISABLE"] = "1"
-# os.environ["NCCL_IB_DISABLE"] = "1"
-# os.environ["MUJOCO_GL"] = "osmesa"
-
-import fire
-from shutil import rmtree
-import gymnasium as gym
-import numpy as np
-
-import torch
-from torch.nn import Module, GRU, Linear
-import torch.nn.functional as F
-
-# functions
-
-def exists(v):
-    return v is not None
-
-def softclamp(t, value):
-    return (t / value).tanh() * value
-
-class HumanoidEnvironment(Module):
-    def __init__(
-        self,
-        video_folder = './recordings_humanoid',
-        render_every_eps = 100,
-        max_steps = 1000,
-        repeats = 1,
-        vectorized = False,
-        num_envs = 1
-    ):
-        super().__init__()
-
-        self.vectorized = vectorized
-        self.num_envs = num_envs
-
-        if vectorized:
-            env = gym.make_vec('Humanoid-v5', num_envs = num_envs, render_mode = 'rgb_array')
-        else:
-            env = gym.make('Humanoid-v5', render_mode = 'rgb_array')
-
-        self.env = env
-        self.max_steps = max_steps
-        self.repeats = repeats
-        self.video_folder = video_folder
-        self.render_every_eps = render_every_eps
-
-    def pre_main_callback(self):
-        # the `pre_main_callback` on the environment passed in is called before the start of the evolutionary strategies loop
-
-        rmtree(self.video_folder, ignore_errors = True)
-
-        if not self.vectorized:
-            self.env = gym.wrappers.RecordVideo(
-                env = self.env,
-                video_folder = self.video_folder,
-                name_prefix = 'recording',
-                episode_trigger = lambda eps_num: (eps_num % self.render_every_eps) == 0,
-                disable_logger = True
-            )
-
-    def forward(self, model):
-
-        device = next(model.parameters()).device
-
-        seed = torch.randint(0, int(1e6), ())
-
-        num_envs = self.num_envs if self.vectorized else 1
-        cum_reward = torch.zeros(num_envs, device = device)
-
-        for _ in range(self.repeats):
-            state, _ = self.env.reset(seed = seed.item())
-
-            step = 0
-            hiddens = None
-            last_action = None
-            
-            dones = torch.zeros(num_envs, device = device, dtype = torch.bool)
-
-            while step < self.max_steps and not dones.all():
-
-                state_torch = torch.from_numpy(state).float().to(device)
-
-                action_logits, hiddens = model(state_torch, hiddens)
-
-                mean, log_var = action_logits.chunk(2, dim = -1)
-
-                # sample and then bound and scale to -0.4 to 0.4
-
-                std = (0.5 * softclamp(log_var, 5.)).exp()
-                sampled = mean + torch.randn_like(mean) * std
-                action = sampled.tanh() * 0.4
-
-                next_state, reward, truncated, terminated, info = self.env.step(action.detach().cpu().numpy() if self.vectorized else action.item())
-
-                reward_np = np.array(reward) if not isinstance(reward, np.ndarray) else reward
-                total_reward_base = torch.from_numpy(reward_np).float().to(device)
-
-                # reward functions
-
-                # encouraged to move forward (1.0) and stay upright (> 1.2 meters)
-
-                z_pos = torch.from_numpy(next_state[..., 0]).float().to(device)
-                x_vel = torch.from_numpy(next_state[..., 5]).float().to(device)
-
-                reward_forward = x_vel
-                reward_upright = (z_pos > 1.2).float()
-
-                exploration_bonus = std.mean(dim = -1) * 0.05
-                penalize_extreme_actions = (mean.abs() > 1.).float().mean(dim = -1) * 0.05
-
-                penalize_action_change = 0.
-                if exists(last_action):
-                    penalize_action_change = (last_action - action).abs().mean(dim = -1) * 0.1
-
-                total_reward = total_reward_base + reward_forward + reward_upright + exploration_bonus - penalize_extreme_actions - penalize_action_change
-
-                # only add reward if not done
-
-                mask = (~dones).float()
-                cum_reward += total_reward * mask
-
-                # update dones
-
-                dones_np = np.array(truncated | terminated) if not isinstance(truncated | terminated, np.ndarray) else (truncated | terminated)
-                dones |= torch.from_numpy(dones_np).to(device)
-
-                step += 1
-
-                state = next_state
-                last_action = action
-
-        if not self.vectorized:
-            return cum_reward.item() / self.repeats
-
-        return cum_reward / self.repeats
-
-# evo strategy
-
-from x_evolution import EvoStrategy
-
-from x_mlps_pytorch.residual_normed_mlp import ResidualNormedMLP
-
-class Model(Module):
-
-    def __init__(self):
-        super().__init__()
-
-        self.deep_mlp = ResidualNormedMLP(
-            dim_in = 348,
-            dim = 256,
-            depth = 8,
-            residual_every = 2
-        )
-
-        self.gru = GRU(256, 256, batch_first = True)
-
-        self.to_pred = Linear(256, 17 * 2, bias = False)
-
-    def forward(self, state, hiddens = None):
-
-        x = self.deep_mlp(state)
-
-        x = x.unsqueeze(-2)
-        gru_out, hiddens = self.gru(x, hiddens)
-        x = x + gru_out
-        x = x.squeeze(-2)
-
-        return self.to_pred(x), hiddens
-
-from torch.optim.lr_scheduler import CosineAnnealingLR
-
-def main(
-    vectorized = False,
-    num_envs = 8
-):
-    evo_strat = EvoStrategy(
-        Model(),
-        environment = HumanoidEnvironment(
-            repeats = 1,
-            render_every_eps = 200,
-            vectorized = vectorized,
-            num_envs = num_envs
-        ),
-        vectorized = vectorized,
-        vector_size = num_envs,
-        num_generations = 50_000,
-        noise_population_size = 200,
-        noise_low_rank = 1,
-        noise_scale = 1e-2,
-        noise_scale_clamp_range = (5e-3, 2e-2),
-        learned_noise_scale = True,
-        use_sigma_optimizer = True,
-        learning_rate = 1e-3,
-        noise_scale_learning_rate = 1e-4,
-        use_scheduler = True,
-        scheduler_klass = CosineAnnealingLR,
-        scheduler_kwargs = dict(T_max = 50_000)
-    )
-
-    evo_strat()
-
-if __name__ == '__main__':
-    fire.Fire(main)

x_evolution-0.1.30/train_lunar.py

@@ -1,141 +0,0 @@
-# /// script
-# dependencies = [
-#     "fire",
-#     "gymnasium[box2d]>=1.0.0",
-#     "gymnasium[other]",
-#     "x-evolution>=0.0.20",
-#     "x-mlps-pytorch>=0.2.0"
-# ]
-# ///
-
-import fire
-from shutil import rmtree
-import gymnasium as gym
-import numpy as np
-
-import torch
-from torch.nn import Module
-import torch.nn.functional as F
-from x_mlps_pytorch.residual_normed_mlp import ResidualNormedMLP
-from torch.optim.lr_scheduler import CosineAnnealingLR
-
-class LunarEnvironment(Module):
-    def __init__(
-        self,
-        video_folder = './recordings',
-        render_every_eps = 500,
-        max_steps = 500,
-        repeats = 1,
-        vectorized = False,
-        num_envs = 1
-    ):
-        super().__init__()
-
-        self.vectorized = vectorized
-        self.num_envs = num_envs
-
-        if vectorized:
-            env = gym.make_vec('LunarLander-v3', num_envs = num_envs, render_mode = 'rgb_array')
-        else:
-            env = gym.make('LunarLander-v3', render_mode = 'rgb_array')
-
-        self.env = env
-        self.max_steps = max_steps
-        self.repeats = repeats
-        self.video_folder = video_folder
-        self.render_every_eps = render_every_eps
-
-    def pre_main_callback(self):
-        # the `pre_main_callback` on the environment passed in is called before the start of the evolutionary strategies loop
-
-        rmtree(self.video_folder, ignore_errors = True)
-
-        if not self.vectorized:
-            self.env = gym.wrappers.RecordVideo(
-                env = self.env,
-                video_folder = self.video_folder,
-                name_prefix = 'recording',
-                episode_trigger = lambda eps_num: (eps_num % self.render_every_eps) == 0,
-                disable_logger = True
-            )
-
-    def forward(self, model):
-
-        device = next(model.parameters()).device
-
-        seed = torch.randint(0, int(1e6), ())
-
-        num_envs = self.num_envs if self.vectorized else 1
-        cum_reward = torch.zeros(num_envs, device = device)
-
-        for _ in range(self.repeats):
-            state, _ = self.env.reset(seed = seed.item())
-
-            step = 0
-            dones = torch.zeros(num_envs, device = device, dtype = torch.bool)
-
-            while step < self.max_steps and not dones.all():
-
-                state_torch = torch.from_numpy(state).to(device)
-
-                action_logits = model(state_torch)
-
-                action = F.gumbel_softmax(action_logits, hard = True).argmax(dim = -1)
-
-                next_state, reward, truncated, terminated, *_ = self.env.step(action.detach().cpu().numpy() if self.vectorized else action.item())
-
-                reward_np = np.array(reward) if not isinstance(reward, np.ndarray) else reward
-                total_reward = torch.from_numpy(reward_np).float().to(device)
-
-                mask = (~dones).float()
-                cum_reward += total_reward * mask
-
-                dones_np = np.array(truncated | terminated) if not isinstance(truncated | terminated, np.ndarray) else (truncated | terminated)
-                dones |= torch.from_numpy(dones_np).to(device)
-
-                step += 1
-
-                state = next_state
-
-        if not self.vectorized:
-            return cum_reward.item() / self.repeats
-
-        return cum_reward / self.repeats
-
-# evo strategy
-
-from x_evolution import EvoStrategy
-
-def main(
-    vectorized = False,
-    num_envs = 8
-):
-    actor = ResidualNormedMLP(dim_in = 8, dim = 24, depth = 2, residual_every = 1, dim_out = 4)
-
-    evo_strat = EvoStrategy(
-        actor,
-        environment = LunarEnvironment(
-            repeats = 2,
-            vectorized = vectorized,
-            num_envs = num_envs
-        ),
-        vectorized = vectorized,
-        vector_size = num_envs,
-        num_generations = 50_000,
-        noise_population_size = 50,
-        noise_low_rank = 1,
-        noise_scale = 1e-2,
-        noise_scale_clamp_range = (5e-3, 2e-2),
-        learned_noise_scale = True,
-        use_sigma_optimizer = True,
-        learning_rate = 1e-3,
-        noise_scale_learning_rate = 1e-4,
-        use_scheduler = True,
-        scheduler_klass = CosineAnnealingLR,
-        scheduler_kwargs = dict(T_max = 50_000)
-    )
-
-    evo_strat()
-
-if __name__ == '__main__':
-    fire.Fire(main)

x_evolution-0.1.30/train_mnist.py

@@ -1,91 +0,0 @@
-# /// script
-# dependencies = [
-#     "fire",
-#     "torchvision",
-#     "x-mlps-pytorch>=0.2.0",
-#     "x-evolution>=0.0.20"
-# ]
-# ///
-
-import fire
-import torch
-from torch import nn
-import torch.nn.functional as F
-from torchvision import datasets, transforms
-from torch.utils.data import DataLoader
-
-# model
-
-from x_mlps_pytorch.residual_normed_mlp import ResidualNormedMLP
-
-model = nn.Sequential(
-    nn.Flatten(),
-    ResidualNormedMLP(dim_in = 784, dim = 512, depth = 8, residual_every = 2, dim_out = 10)
-).half()
-
-batch_size = 256
-
-# data
-
-dataset = datasets.MNIST('./data', train = True, download = True, transform = transforms.ToTensor())
-
-# fitness as inverse of loss
-
-def mnist_environment(
-    model,
-    num_envs = 1,
-    vectorized = False,
-    batch_size = 256
-):
-    device = next(model.parameters()).device
-    
-    iters = num_envs if vectorized else 1
-
-    losses = []
-
-    for _ in range(iters):
-        dataloader = DataLoader(dataset, batch_size = batch_size, shuffle = True)
-        data_iterator = iter(dataloader)
-        data, target = next(data_iterator)
-
-        data, target = data.to(device), target.to(device)
-
-        with torch.no_grad():
-            logits = model(data.half())
-            loss = F.cross_entropy(logits, target)
-        
-        losses.append(-loss)
-
-    if not vectorized:
-        return losses[0]
-
-    return torch.stack(losses)
-
-# evo
-
-from x_evolution import EvoStrategy
-
-def main(
-    vectorized = False,
-    num_envs = 8,
-    batch_size = 256
-):
-    evo_strat = EvoStrategy(
-        model,
-        environment = lambda model: mnist_environment(model, num_envs = num_envs, vectorized = vectorized, batch_size = batch_size),
-        vectorized = vectorized,
-        vector_size = num_envs,
-        noise_population_size = 100,
-        noise_scale = 1e-2,
-        noise_scale_clamp_range = (8e-3, 2e-2),
-        noise_low_rank = 1,
-        num_generations = 10_000,
-        learning_rate = 1e-3,
-        learned_noise_scale = True,
-        noise_scale_learning_rate = 2e-5
-    )
-
-    evo_strat()
-
-if __name__ == '__main__':
-    fire.Fire(main)

x_evolution-0.1.30/train_xor.py

@@ -1,83 +0,0 @@
-import fire
-import torch
-from torch import nn
-import torch.nn.functional as F
-from torch.optim.lr_scheduler import LambdaLR
-
-# model
-
-from torch import nn
-
-model = nn.Sequential(
-    nn.Linear(2, 16),
-    nn.ReLU(),
-    nn.Linear(16, 2)
-).half()
-
-batch_size = 128
-
-# fitness as inverse of loss
-
-from x_evolution import EvoStrategy
-
-def xor_environment(
-    model,
-    num_envs = 1,
-    vectorized = False,
-    batch_size = 128
-):
-    device = next(model.parameters()).device
-
-    iters = num_envs if vectorized else 1
-
-    losses = []
-
-    for _ in range(iters):
-        data = torch.randint(0, 2, (batch_size, 2))
-        labels = data[:, 0] ^ data[:, 1]
-
-        data, labels = tuple(t.to(device) for t in (data, labels))
-
-        with torch.no_grad():
-            logits = model(data.half())
-            loss = F.cross_entropy(logits, labels)
-        
-        losses.append(-loss)
-
-    if not vectorized:
-        return losses[0]
-
-    return torch.stack(losses)
-
-# evo
-
-def main(
-    vectorized = False,
-    num_envs = 8,
-    batch_size = 128
-):
-    evo_strat = EvoStrategy(
-        model,
-        environment = lambda model: xor_environment(model, num_envs = num_envs, vectorized = vectorized, batch_size = batch_size),
-        vectorized = vectorized,
-        vector_size = num_envs,
-        noise_population_size = 100,
-        noise_low_rank = 1,
-        num_generations = 100_000,
-        learning_rate = 1e-1,
-        noise_scale = 1e-1,
-        noise_scale_clamp_range = (0.05, 0.2),
-        learned_noise_scale = True,
-        noise_scale_learning_rate = 5e-4,
-        use_scheduler = True,
-        scheduler_klass = LambdaLR,
-        scheduler_kwargs = dict(lr_lambda = lambda step: min(1., step / 10.)),
-        use_sigma_scheduler = True,
-        sigma_scheduler_klass = LambdaLR,
-        sigma_scheduler_kwargs = dict(lr_lambda = lambda step: min(1., step / 10.))
-    )
-
-    evo_strat()
-
-if __name__ == '__main__':
-    fire.Fire(main)