QDax：加速质量多样性优化框架

颜妙瑶Titus

于 2024-09-11 09:50:39 发布

阅读量747

点赞数 17

本文链接：https://blog.csdn.net/gitblog_00636/article/details/142129668

版权

QDax：加速质量多样性优化框架

QDax Accelerated Quality-Diversity 项目地址: https://gitcode.com/gh_mirrors/qd/QDax

项目介绍

QDax 是一个基于 Jax 的开源库，专注于加速质量多样性（Quality-Diversity, QD）算法与神经进化计算。该框架设计灵活、易于扩展，适用于各种优化场景，从黑盒优化到连续控制任务不等。QDax 不仅支持即时编译以提高效率，还特别强调硬件加速，使得以往需在大型CPU集群上运行数天至数周的QD算法能够在分钟内完成。它包含了多种流行的QD、神经进化及强化学习算法实现，并提供丰富示例。

项目快速启动

要迅速上手并运行QDax中的一个简单示例，首先确保你的环境已安装Jax以及对应的依赖。然后，通过以下命令安装QDax：

pip install qdax

如果你的系统支持CUDA且想启用CUDA 12的支持，可以使用以下命令安装：

pip install qdax[cuda12]

接下来，一个简单的例子展示了如何利用QDax进行基本的MAP-Elites算法设置：

import jax.random
from qdax.core.map_elites import MAPElites
from qdax.core.containers.mapelites_repertoire import compute_euclidean_centroids
from qdax.tasks.arm import arm_scoring_function
from qdax.core.emitters.mutation_operators import isoline_variation
from qdax.core.emitters.standard_emitters import MixingEmitter
from qdax.utils.metrics import default_qd_metrics

seed = 42
num_param_dimensions = 100
arm_init_batch_size = 100
batch_size = 1024
num_iterations = 50
grid_shape = (100, 100)
min_param = 0.0
max_param = 1.0
min_bd = 0.0
max_bd = 1.0

# 初始化随机种子和控制器群体
random_key = jax.random.PRNGKey(seed)
subkey = jax.random.split(random_key)[0]
init_variables = jax.random.uniform(subkey, shape=(arm_init_batch_size, num_param_dimensions), minval=min_param)

# 运行MAP-Elites算法
for i in range(num_iterations):
    repertoire, emitter_state, metrics, random_key = MAPElites.update(
        repertoire, emitter_state, random_key
    )