57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import os
|
|
os.environ['MKL_THREADING_LAYER'] = 'GNU'
|
|
os.environ["OMP_NUM_THREADS"] = "8"
|
|
os.environ["MKL_NUM_THREADS"] = "8"
|
|
import torch
|
|
import numpy as np
|
|
from torch.nn import Module
|
|
import gymnasium as gym
|
|
from gym_env import CustomEnv
|
|
from stable_baselines3 import PPO
|
|
from stable_baselines3.common.vec_env import SubprocVecEnv
|
|
from stable_baselines3.common.vec_env import DummyVecEnv
|
|
from sb3_contrib import RecurrentPPO
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
|
|
current_dir = os.path.dirname(os.path.abspath("__file__"))
|
|
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
|
|
|
|
class Sin(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x):
|
|
return torch.sin(x)
|
|
|
|
if __name__ == '__main__':
|
|
|
|
vec_env = CustomEnv(device_id=1)
|
|
name = "d1a3o12_c1"
|
|
|
|
model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_a0"), env=vec_env, device=torch.device("cuda:1"))
|
|
|
|
# model = PPO(
|
|
# "MlpPolicy",
|
|
# policy_kwargs=dict(activation_fn=Sin),
|
|
# env=vec_env,
|
|
# device=torch.device("cuda:1"),
|
|
# verbose=0)
|
|
|
|
writer = SummaryWriter(log_dir=os.path.join(parent_dir, "tensorboard", name))
|
|
max_reward = 0
|
|
|
|
for i in range(100):
|
|
model.learn(total_timesteps=480)
|
|
test_env = model.get_env()
|
|
test_obs = test_env.reset()
|
|
list_reward = []
|
|
for step in range(480):
|
|
test_action, _states = model.predict(observation=test_obs)
|
|
test_obs, test_rewards, test_dones, info = test_env.step(test_action)
|
|
list_reward.append(test_rewards)
|
|
|
|
avg_reward = np.mean(list_reward[-240:])
|
|
writer.add_scalar('Reward', np.mean(avg_reward), i)
|
|
if avg_reward > max_reward:
|
|
max_reward = avg_reward
|
|
model.save(os.path.join(parent_dir, "models", name + ".zip")) |