Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions experiments/ppo_4x4grid.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@


if __name__ == "__main__":
print(os.getcwd())
ray.init()

env_name = "4x4grid"
Expand Down
30 changes: 6 additions & 24 deletions experiments/sb3_grid4x4.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
import numpy as np
import supersuit as ss
import traci
from pyvirtualdisplay.smartdisplay import SmartDisplay
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import EvalCallback
from stable_baselines3.common.evaluation import evaluate_policy
Expand All @@ -16,34 +15,26 @@


if __name__ == "__main__":
RESOLUTION = (3200, 1800)

env = sumo_rl.grid4x4(use_gui=True, out_csv_name="outputs/grid4x4/ppo_test", virtual_display=RESOLUTION)
env = sumo_rl.grid4x4(use_gui=False, out_csv_name="outputs/grid4x4/ppo_train")

max_time = env.unwrapped.env.sim_max_time
delta_time = env.unwrapped.env.delta_time

print("Environment created")

env = ss.pettingzoo_env_to_vec_env_v1(env)
env = ss.concat_vec_envs_v1(env, 2, num_cpus=1, base_class="stable_baselines3")
env = ss.concat_vec_envs_v1(env, 2, num_cpus=16, base_class="stable_baselines3")
env = VecMonitor(env)

model = PPO(
"MlpPolicy",
env,
verbose=3,
gamma=0.95,
n_steps=256,
ent_coef=0.0905168,
learning_rate=0.00062211,
vf_coef=0.042202,
max_grad_norm=0.9,
gae_lambda=0.99,
n_epochs=5,
clip_range=0.3,
batch_size=256,
tensorboard_log="./logs/grid4x4/ppo_test",
tensorboard_log="./logs/grid4x4/ppo_train",
)

print("Starting training")
Expand All @@ -55,28 +46,19 @@
print(mean_reward)
print(std_reward)

model.save('ppo_output')

# Maximum number of steps before reset, +1 because I'm scared of OBOE
print("Starting rendering")
num_steps = (max_time // delta_time) + 1

obs = env.reset()

if os.path.exists("temp"):
shutil.rmtree("temp")

os.mkdir("temp")
# img = disp.grab()
# img.save(f"temp/img0.jpg")

img = env.render()
for t in trange(num_steps):
actions, _ = model.predict(obs, state=None, deterministic=False)
obs, reward, done, info = env.step(actions)
img = env.render()
img.save(f"temp/img{t}.jpg")

subprocess.run(["ffmpeg", "-y", "-framerate", "5", "-i", "temp/img%d.jpg", "output.mp4"])
env.render()

print("All done, cleaning up")
shutil.rmtree("temp")
env.close()