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()
[强化学习]-- highway_env 自动驾驶仿真
最新推荐文章于 2024-08-17 21:49:24 发布