gym实时渲染模型训练效果

深度强化学习模型在训练玩游戏时，不渲染游戏画面可以大大加快训练速度

# 将 render_mode 设置为 None 就不会渲染游戏，设置为 "human" 就会以图像帧的形式渲染
# 注：老版本需要调用 render 方法才会渲染，新版本会在调用 reset 方法时就自动渲染
env = gymnasium.make("CartPole-v1", render_mode='None')

如果我想在训练过程中实时查看模型的训练效果，则可以使用多线程实现：

一个线程负责训练模型，另一个线程中重新创建一个一样的游戏环境，负责渲染，它们都从主线程获取最新的模型（训练线程就可以作为主线程）

import threading

def render(model):
    render_env = gymnasium.make("CartPole-v1", render_mode='human')
    with torch.no_grad():
        while True:
            obs, _ = render_env.reset()
            while True:
                act = convert_model_output_to_act(model(torch.FloatTensor(obs)))
                _, _, terminated, _, _ = render_env.step(act)
                if terminated: break

m = MyPolicy()

# 创建一个渲染线程，并将正在训练的模型作为参数传递进去，daemon=True 表示将其作为守护线程，主线程退出后它也跟着退出
threading.Thread(target=render, args=(m,), daemon=True).start()

# 主线程就可以作为训练线程，下面部分就是模型训练相关的代码
train_env = gymnasium.make("CartPole-v1", render_mode='None')
...

或者使用子进程也可以，使用方法基本一致，但是要注意，1）这里的 main 方法是必须的，2）进程中传递的args参数并不是引用传递，而是值传递，意味着在主进程中修改了actor的参数后，子进程仍然使用的是原来的actor：

import multiprocessing
...
if __name__ == "__main__":
    multiprocessing.Process(target=visualize_train, args=(actor,), daemon=True).start()

Leave a Comment Cancel reply