[强化学习]-- highway_env 自动驾驶仿真

import os
from stable_baselines3 import DQN
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv, VecVideoRecorder
from stable_baselines3.common.callbacks import EvalCallback
from highway_env import envs

# 创建环境
env_id = "highway-v0"
env = make_vec_env(env_id, n_envs=1)

# 训练参数
total_timesteps = 500000
log_dir = "logs/"
os.makedirs(log_dir, exist_ok=True)

# 创建模型
model = DQN('MlpPolicy', env, verbose=1, tensorboard_log=log_dir)

# 设置评估回调
eval_callback = EvalCallback(env, best_model_save_path=log_dir,
                             log_path=log_dir, eval_freq=5000,
                             deterministic=True, render=False)

# 开始训练
model.learn(total_timesteps=total_timesteps, callback=eval_callback)

# 保存模型
model.save("dqn_highway")

# 评估模型
mean_reward, std_reward = evaluate_policy(model, env, n_eval_episodes=10)
print(f"Mean reward: {mean_reward:.2f} +/- {std_reward:.2f}")

# 渲染模型行为
obs = env.reset()
for _ in range(1000):
    action, _states = model.predict(obs, deterministic=True)
    obs, rewards, dones, info = env.step(action)
    env.render()
env.close()