EVO X
Open menu

在 EvoX 中求解 Brax 问题

在 EvoX 中求解 Brax 问题

EvoX 深度结合 Brax 进行神经进化研究。 这里我们将展示一个在 EvoX 中求解 Brax 问题的示例。

# install EvoX and Brax, skip it if you have already installed EvoX or Brax
from importlib.util import find_spec
from IPython.display import HTML

if find_spec("evox") is None:
    %pip install evox
if find_spec("brax") is None:
    %pip install brax
# The dependent packages or functions in this example
import torch
import torch.nn as nn

from evox.algorithms import PSO
from evox.problems.neuroevolution.brax import BraxProblem
from evox.utils import ParamsAndVector
from evox.workflows import EvalMonitor, StdWorkflow

什么是 Brax

Brax 是一个快速且完全可微的物理引擎，用于机器人、人类感知、材料科学、强化学习以及其他重度依赖模拟的应用的研究与开发。

这里我们将演示 Brax 的 “swimmer” 环境。

欲了解更多信息，您可以浏览 Brax 的 Github。

设计神经网络类

首先，我们需要决定构建哪种神经网络。

这里我们将给出一个简单的多层感知机 (MLP) 类。

# Construct an MLP using PyTorch.
# This MLP has 3 layers.


class SimpleMLP(nn.Module):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        self.features = nn.Sequential(nn.Linear(17, 8), nn.Tanh(), nn.Linear(8, 6))

    def forward(self, x):
        x = self.features(x)
        return torch.tanh(x)

初始化模型

通过 SimpleMLP 类，我们可以初始化一个 MLP 模型。

# Make sure that the model is on the same device, better to be on the GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
# Reset the random seed
seed = 1234
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

# Initialize the MLP model
model = SimpleMLP().to(device)

初始化适配器

适配器可以帮助我们进行数据的来回转换。

adapter = ParamsAndVector(dummy_model=model)

有了适配器，我们就可以着手进行这项神经进化任务了。

设置运行流程

初始化算法和问题

我们初始化一个 PSO 算法，问题是在 “swimmer” 环境下的 Brax 问题。

# Set the population size
POP_SIZE = 1024

# Get the bound of the PSO algorithm
model_params = dict(model.named_parameters())
pop_center = adapter.to_vector(model_params)
lower_bound = torch.full_like(pop_center, -5)
upper_bound = torch.full_like(pop_center, 5)

# Initialize the PSO, and you can also use any other algorithms
algorithm = PSO(
    pop_size=POP_SIZE,
    lb=lower_bound,
    ub=upper_bound,
    device=device,
)

# Initialize the Brax problem
problem = BraxProblem(
    policy=model,
    env_name="halfcheetah",
    max_episode_length=1000,
    num_episodes=3,
    pop_size=POP_SIZE,
    device=device,
)

在本例中，我们将为每个回合（episode）使用 1000 步，并返回 3 个回合的平均奖励作为适应度值。

设置监控器

# set an monitor, and it can record the top 3 best fitnesses
monitor = EvalMonitor(
    topk=3,
    device=device,
)

初始化工作流

# Initiate an workflow
workflow = StdWorkflow(
    algorithm=algorithm,
    problem=problem,
    monitor=monitor,
    opt_direction="max",
    solution_transform=adapter,
    device=device,
)

运行工作流

运行工作流，见证奇迹！

注意： 下面的代码块大约需要运行 20 分钟。 具体时间可能因您的硬件而异。

# Set the maximum number of generations
max_generation = 50

# Run the workflow
workflow.init_step()
compiled_step = torch.compile(workflow.step)
for i in range(max_generation):
    if i % 10 == 0:
        print(f"Generation {i}")
    compiled_step()

print(f"Top fitness: {monitor.get_best_fitness()}")
best_params = adapter.to_params(monitor.get_best_solution())
print(f"Best params: {best_params}")
monitor.get_best_fitness()
monitor.plot()
html_string = problem.visualize(best_params)
escaped_string = html_string.replace('"', """)
HTML(f'<iframe srcdoc="{escaped_string}" width="100%" height="480" frameborder="0"></iframe>')

重要提示：

  • 通常情况下，您只需要 HTML(problem.visualize(best_params)) 即可进行渲染。上面的代码是一个变通方法，以确保结果在我们的网站上正确显示。
  • PSO 算法并未针对此类任务进行专门优化，因此性能可能会受到限制。本示例仅用于演示目的。

我们希望您享受使用 EvoX 求解 Brax 问题的过程，祝您玩得开心！

← 返回文档首页