x-evolution 0.1.25__tar.gz → 0.1.27__tar.gz

{x_evolution-0.1.25 → x_evolution-0.1.27}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.25
+Version: 0.1.27
 Summary: x-evolution
 Project-URL: Homepage, https://pypi.org/project/x-evolution/
 Project-URL: Repository, https://github.com/lucidrains/x-evolution
@@ -38,7 +38,7 @@ Requires-Dist: accelerate
 Requires-Dist: beartype
 Requires-Dist: einops>=0.8.0
 Requires-Dist: torch>=2.4
-Requires-Dist: x-mlps-pytorch>=0.1.31
+Requires-Dist: x-mlps-pytorch>=0.2.0
 Requires-Dist: x-transformers>=2.11.23
 Provides-Extra: examples
 Requires-Dist: gymnasium[box2d]>=1.0.0; extra == 'examples'

{x_evolution-0.1.25 → x_evolution-0.1.27}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-evolution"
-version = "0.1.25"
+version = "0.1.27"
 description = "x-evolution"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -27,7 +27,7 @@ dependencies = [
     "beartype",
     "einops>=0.8.0",
     "torch>=2.4",
-    "x-mlps-pytorch>=0.1.31",
+    "x-mlps-pytorch>=0.2.0",
     "x-transformers>=2.11.23"
 ]
 

{x_evolution-0.1.25 → x_evolution-0.1.27}/tests/test_evolution.py

@@ -47,3 +47,24 @@ def test_evo_strat(
     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.27/train_humanoid.py

@@ -0,0 +1,215 @@
+# /// 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.27/train_lunar.py

@@ -0,0 +1,141 @@
+# /// 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.27/train_mnist.py

@@ -0,0 +1,91 @@
+# /// 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.27/train_xor.py

@@ -0,0 +1,83 @@
+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)

{x_evolution-0.1.25 → x_evolution-0.1.27}/x_evolution/x_evolution.py

@@ -11,6 +11,7 @@ from torch.nn import Module, ModuleList, Parameter, ParameterList
 from torch.optim import SGD, Adam, Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
+import torch.distributed as dist
 import torch.nn.functional as F
 
 from beartype import beartype
@@ -93,11 +94,17 @@ class EvoStrategy(Module):
         verbose = True,
         accelerator: Accelerator | None = None,
         accelerate_kwargs: dict = dict(),
-        reject_generation_fitnesses_if: Callable[[Tensor], bool] | None = None
+        reject_generation_fitnesses_if: Callable[[Tensor], bool] | None = None,
+        vectorized = False,
+        vector_size: int | None = None,
+        sync_on_init = True
     ):
         super().__init__()
         self.verbose = verbose
 
+        self.vectorized = vectorized
+        self.vector_size = vector_size
+
         if not exists(accelerator):
             accelerator = Accelerator(cpu = cpu, **accelerate_kwargs)
 
@@ -122,12 +129,10 @@ class EvoStrategy(Module):
         self.model = model
         self.noisable_model = Noisable(model, low_rank = noise_low_rank)
 
-        # use prepare and run through environment once to sync params
-
-        wrapped_model = accelerator.prepare(model)
+        # maybe sync model params and buffers
 
-        with torch.no_grad():
-            environment(wrapped_model)
+        if sync_on_init:
+            self.sync_model_params_and_buffers_()
 
         # get param dictionary
 
@@ -249,6 +254,17 @@ class EvoStrategy(Module):
     def device(self):
         return self.accelerate.device
 
+    @torch.no_grad()
+    def sync_model_params_and_buffers_(self):
+        if not self.accelerate.num_processes > 1:
+            return
+
+        for param in self.model.parameters():
+            dist.broadcast(param, src = 0)
+
+        for buffer in self.model.buffers():
+            dist.broadcast(buffer, src = 0)
+
     def print(self, *args, **kwargs):
         if not self.verbose:
             return
@@ -475,24 +491,28 @@ class EvoStrategy(Module):
                     fitnesses.append([0., 0.] if self.mirror_sampling else 0.)
                     continue
 
-                individual_param_seeds = with_seed(individual_seed)(randint)(0, MAX_SEED_VALUE, (self.num_params,))
-
-                noise_config = dict(zip(self.param_names_to_optimize, individual_param_seeds.tolist()))
-
-                # determine noise scale, which can be fixed or learned
+                def get_fitness(negate = False):
+                    individual_param_seeds = with_seed(individual_seed.item())(randint)(0, MAX_SEED_VALUE, (self.num_params,))
+                    noise_config = dict(zip(self.param_names_to_optimize, individual_param_seeds.tolist()))
 
-                noise_config_with_scale = dict()
+                    noise_config_with_scale = dict()
+                    for param_name, seed in noise_config.items():
+                        noise_scale = self._get_noise_scale(param_name)
+                        noise_config_with_scale[param_name] = (seed, noise_scale)
 
-                for param_name, seed in noise_config.items():
+                    with model.temp_add_noise_(noise_config_with_scale, negate = negate):
+                        fitness = with_seed(maybe_rollout_seed)(self.environment)(model)
 
-                    noise_scale = self._get_noise_scale(param_name)
+                    if isinstance(fitness, Tensor) and fitness.numel() > 1:
+                        fitness = fitness.mean().item()
+                    elif isinstance(fitness, Tensor):
+                        fitness = fitness.item()
 
-                    noise_config_with_scale[param_name] = (seed, noise_scale)
+                    return fitness
 
-                # maybe roll out with a fixed seed
+                # evaluate
 
-                with model.temp_add_noise_(noise_config_with_scale):
-                    fitness = with_seed(maybe_rollout_seed)(rollout_for_fitness)()
+                fitness = get_fitness(negate = False)
 
                 if not self.mirror_sampling:
                     fitnesses.append(fitness)
@@ -500,8 +520,7 @@ class EvoStrategy(Module):
 
                 # handle mirror sampling
 
-                with model.temp_add_noise_(noise_config_with_scale, negate = True):
-                    fitness_mirrored = with_seed(maybe_rollout_seed)(rollout_for_fitness)()
+                fitness_mirrored = get_fitness(negate = True)
 
                 fitnesses.append([fitness, fitness_mirrored])
 

x_evolution-0.1.25/train_humanoid.py

@@ -1,130 +0,0 @@
-# /// script
-# dependencies = [
-#     "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"
-
-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
-
-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
-    ):
-        super().__init__()
-
-        # Humanoid-v5
-        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)
-
-        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), ())
-
-        cum_reward = 0.
-
-        for _ in range(self.repeats):
-            state, _ = self.env.reset(seed = seed.item())
-
-            step = 0
-
-            while step < self.max_steps:
-
-                state = torch.from_numpy(state).float().to(device)
-
-                action_logits = model(state)
-
-                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, *_ = self.env.step(action.detach().cpu().numpy())
-
-                cum_reward += float(reward)
-                step += 1
-
-                state = next_state
-
-                if truncated or terminated:
-                    break
-
-        return cum_reward / self.repeats
-
-# evo strategy
-
-from x_evolution import EvoStrategy
-
-from x_mlps_pytorch.residual_normed_mlp import ResidualNormedMLP
-
-actor = ResidualNormedMLP(
-    dim_in = 348, # state
-    dim = 256,
-    depth = 8,
-    residual_every = 2,
-    dim_out = 17 * 2 # action mean logvar
-)
-
-from torch.optim.lr_scheduler import CosineAnnealingLR
-
-evo_strat = EvoStrategy(
-    actor,
-    environment = HumanoidEnvironment(repeats = 2),
-    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()

x_evolution-0.1.25/train_lunar.py

@@ -1,107 +0,0 @@
-# /// script
-# dependencies = [
-#     "gymnasium[box2d]>=1.0.0",
-#     "gymnasium[other]",
-#     "x-evolution>=0.0.20"
-# ]
-# ///
-
-from shutil import rmtree
-import gymnasium as gym
-
-import torch
-from torch.nn import Module
-import torch.nn.functional as F
-
-class LunarEnvironment(Module):
-    def __init__(
-        self,
-        video_folder = './recordings',
-        render_every_eps = 500,
-        max_steps = 500,
-        repeats = 1
-    ):
-        super().__init__()
-
-        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)
-
-        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), ())
-
-        cum_reward = 0.
-
-        for _ in range(self.repeats):
-            state, _ = self.env.reset(seed = seed.item())
-
-            step = 0
-
-            while step < self.max_steps:
-
-                state = torch.from_numpy(state).to(device)
-
-                action_logits = model(state)
-
-                action = F.gumbel_softmax(action_logits, hard = True).argmax(dim = -1)
-
-                next_state, reward, truncated, terminated, *_ = self.env.step(action.item())
-
-                cum_reward += float(reward)
-                step += 1
-
-                state = next_state
-
-                if truncated or terminated:
-                    break
-
-        return cum_reward / self.repeats
-
-# evo strategy
-
-from x_evolution import EvoStrategy
-
-from x_mlps_pytorch.residual_normed_mlp import ResidualNormedMLP
-
-actor = ResidualNormedMLP(dim_in = 8, dim = 24, depth = 2, residual_every = 1, dim_out = 4)
-
-from torch.optim.lr_scheduler import CosineAnnealingLR
-
-evo_strat = EvoStrategy(
-    actor,
-    environment = LunarEnvironment(repeats = 2),
-    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()

x_evolution-0.1.25/train_mnist.py

@@ -1,63 +0,0 @@
-# /// script
-# dependencies = [
-#     "torchvision",
-#     "x-evolution>=0.0.20"
-# ]
-# ///
-
-import torch
-from torch import tensor, 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 loss_mnist(model):
-    device = next(model.parameters()).device
-    
-    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)
-
-    return -loss
-
-# evo
-
-from x_evolution import EvoStrategy
-
-evo_strat = EvoStrategy(
-    model,
-    environment = loss_mnist,
-    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()

x_evolution-0.1.25/train_xor.py

@@ -1,57 +0,0 @@
-import torch
-from torch import tensor
-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 loss_xor(model):
-    device = next(model.parameters()).device
-
-    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)
-
-    return -loss
-
-# evo
-
-evo_strat = EvoStrategy(
-    model,
-    environment = loss_xor,
-    noise_population_size = 100,
-    noise_low_rank = 1,
-    num_generations = 100_000,
-    learning_rate = 1e-1,
-    noise_scale = 1e-1,
-    noise_scale_clamp_range = (5e-2, 2e-1),
-    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()