evox.problems.neuroevolution.mujoco_playground
¶
模块内容¶
类¶
MJX 问题封装器。 |
API¶
- class evox.problems.neuroevolution.mujoco_playground.MujocoProblem(policy: torch.nn.Module, env_name: str, max_episode_length: int, num_episodes: int, seed: int = None, pop_size: int | None = None, rotate_key: bool = True, reduce_fn: Callable[[torch.Tensor, int], torch.Tensor] = torch.mean, device: torch.device | None = None)¶
Bases:
evox.core.Problem
MJX 问题封装器。
初始化
构建一个基于MujocoPlayground的问题。首先,你需要定义一个policy模型。然后,你需要设置
environment name <https://github.com/google-deepmind/mujoco_playground/tree/main/mujoco_playground/_src/registry.py>
,最大episode长度,以及评估每个个体的episode数量。对于每个个体,将使用不同的seed在环境中运行该policy num_episodes次,并使用reduce_fn来合并奖励(默认是平均值)。在每次迭代中,不同个体将共享同一组随机键。- 参数:
policy -- 该策略模型的前向函数是: `forward(batched_obs) -> action
env_name -- 环境名。
max_episode_length -- 每个episode的最大时间步数。
num_episodes -- 用于评估每个个体的幕数量。
seed -- 用于为Mujoco环境创建PRNGKey的种子。当为None时,随机选择一个。默认值为None。
pop_size -- 要评估的种群大小。如果为 None,我们期望输入的种群大小为 1。
rotate_key -- 表示是否在每次迭代中旋转随机键(默认为 True)。
如果为 True,随机键将在每次迭代后旋转,从而导致非确定性且可能有噪声的适应度评估。这意味着相同的策略权重在不同迭代中可能会产生不同的适应度值。
如果为 False,随机键在所有迭代中保持不变,以确保适应度评估的一致性。reduce_fn -- 用于减少多个episode奖励的函数。默认值为torch.mean。
device -- 用于运行计算的设备。默认为当前默认设备。
示例
from evox import problems problem = problems.neuroevolution.MujocoProblem( env_name="SwimmerSwimmer6", policy=model, max_episode_length=1000, num_episodes=3, pop_size=100, rotate_key=False, )
备注
初始密钥是从
torch.random.get_rng_state()
获得的。警告
此问题不支持 HPO 包装器 (
problems.hpo_wrapper.HPOProblemWrapper
) 默认使用,即包含此问题的工作流无法被 vmapped然而,通过将pop_size
设置为内种群大小与外种群大小的乘积,仍然可以在 HPO 工作流中使用此问题。但 HPO 包装器的num_repeats
必须 设置为 1,请使用参数num_episodes
替代。- _evaluate_mjx_record(key: torch.Tensor, model_state: Dict[str, torch.Tensor]) Tuple[Dict[str, torch.Tensor], torch.Tensor, List[Any]] ¶
- evaluate(pop_params: Dict[str, torch.nn.Parameter]) torch.Tensor ¶
评估一组模型参数种群(批量)的最终奖励。
- 参数:
pop_params -- 参数字典,其中每个键是参数名称,每个值是形状为 (batch_size, *param_shape) 的张量,表示批量模型的批量参数。
- 返回:
形状为 (batch_size,) 的张量,包含种群中每个样本的奖励。
- visualize(weights: Dict[str, torch.nn.Parameter], seed: int | None = None, output_type: Literal[mp4, gif] = 'mp4', output_path: str = 'output_video', /, **kwargs) str ¶
可视化带有指定策略和权重的brax环境。
- 参数:
weights -- 策略模型的权重。这是一个参数字典。
seed -- 用于为 brax 环境创建 PRNGKey 的种子。当为 None 时,使用当前的 key。默认值为 None。
output_type -- 可视化的输出类型,“mp4”或“gif”。默认值为“mp4”。
output_path -- 保存输出视频的路径。默认值为 "output_video"。
kwargs -- 传递给 MjxEnv.render() 函数的额外参数。
- 返回:
保存文件的路径。