如何避免多进程复制无关的MLAgents-Unity连接对象？

问题背景 #

我正在实现NEAT（神经网络拓扑进化算法），由于前馈网络评估是训练过程中的性能瓶颈，打算用多进程并行运行。我使用MLAgents库连接Unity，通过仿真评估适应度。

遇到的核心问题：创建进程时，全局定义的MLAgents环境连接对象会随内存栈被复制到每个子进程，导致原连接被关闭，但该连接对象与多进程执行的任务完全无关。如何防止该连接对象被包含到各个子进程中？

已尝试：将多进程代码拆分到单独文件，但问题依旧。

相关代码 #

多进程逻辑代码：

from multiprocessing import Queue
import multiprocessing as mp


def get_action(network, obs, agent_num, queue):
    queue.put([agent_num, network.activate(obs)])

def get_actions(policies,
                fixed_policy,
                fixed_opponent,
                nn_input,
                decision_steps_blue,
                decision_steps_purple,
                agent_count,
                local_to_agent_map):
    # Concurrency things
    num_workers = mp.cpu_count()
    print("CPU Cores: " + str(num_workers))
    pool = mp.Pool(processes=num_workers)  # Problem: Unity connection (MLAgents) being duped
    q = Queue()

    for agent in range(agent_count):
        if local_to_agent_map[agent] in decision_steps_purple or local_to_agent_map[agent] in decision_steps_blue:
            if local_to_agent_map[agent] in decision_steps_blue or not fixed_opponent:
                policy = policies[agent]
            elif fixed_opponent:
                policy = fixed_policy
            pool.apply_async(get_action, args=(policy, nn_input[agent], agent, q))

    pool.close()
    pool.join()

    return q

全局定义的Unity连接对象：

from mlagents_envs.environment import UnityEnvironment

env = UnityEnvironment(seed=1, side_channels=[])   # Object to avoid duplicating

解决方案 #

核心思路是阻止子进程继承或持有不必要的Unity连接对象，以下是几种可行方案：

1. 延迟初始化Unity连接对象 #

不在全局作用域创建env对象，仅在主进程的执行逻辑中初始化，确保子进程启动时全局作用域无该对象：

from mlagents_envs.environment import UnityEnvironment

def init_unity_env():
    global env
    env = UnityEnvironment(seed=1, side_channels=[])

# 主进程入口
if __name__ == "__main__":
    init_unity_env()
    # 后续调用get_actions等多进程逻辑

2. 子进程启动时主动清理全局对象 #

利用multiprocessing.Pool的initializer参数，在每个子进程初始化时清除全局的Unity连接对象：

def clear_unity_env():
    global env
    env = None  # 子进程不需要该连接，直接置空

# 修改Pool创建逻辑
pool = mp.Pool(processes=num_workers, initializer=clear_unity_env)

3. 改用spawn多进程启动方式（推荐） #

Unix系统默认用fork方式启动子进程，会复制父进程全部内存状态；spawn方式会重新启动Python解释器，仅导入必要模块，不复制父进程全局对象：

import multiprocessing as mp

if __name__ == "__main__":
    mp.set_start_method('spawn')  # 设置启动方式
    init_unity_env()  # 初始化Unity环境
    # 后续执行多进程逻辑

额外优化：替换全局队列 #

原代码使用的Queue可通过apply_async的返回结果替代，减少进程间通信开销：

def get_action(network, obs, agent_num):
    return [agent_num, network.activate(obs)]

def get_actions(policies,
                fixed_policy,
                fixed_opponent,
                nn_input,
                decision_steps_blue,
                decision_steps_purple,
                agent_count,
                local_to_agent_map):
    num_workers = mp.cpu_count()
    print(f"CPU Cores: {num_workers}")
    pool = mp.Pool(processes=num_workers, initializer=clear_unity_env)
    results = []

    for agent in range(agent_count):
        agent_id = local_to_agent_map[agent]
        if agent_id in decision_steps_purple or agent_id in decision_steps_blue:
            policy = policies[agent] if (agent_id in decision_steps_blue or not fixed_opponent) else fixed_policy
            res = pool.apply_async(get_action, args=(policy, nn_input[agent], agent))
            results.append(res)

    pool.close()
    pool.join()

    return [res.get() for res in results]

内容的提问来源于stack exchange，提问作者Kristian T