72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
import os
|
|
os.environ['MKL_THREADING_LAYER'] = 'GNU'
|
|
os.environ["OMP_NUM_THREADS"] = "16"
|
|
os.environ["MKL_NUM_THREADS"] = "16"
|
|
import torch
|
|
import numpy as np
|
|
from torch.nn import Module
|
|
import gymnasium as gym
|
|
from gym_env_vortex 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
|
|
import pickle
|
|
|
|
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=3)
|
|
name = "vortex_taylor"
|
|
|
|
model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_re100"), env=vec_env, device=torch.device("cuda:3"))
|
|
|
|
# model = PPO(
|
|
# "MlpPolicy",
|
|
# policy_kwargs=dict(activation_fn=Sin),
|
|
# env=vec_env,
|
|
# device=torch.device("cuda:3"),
|
|
# n_steps=3600,
|
|
# batch_size=360,
|
|
# verbose=0)
|
|
|
|
writer = SummaryWriter(log_dir=os.path.join(parent_dir, "tensorboard", name))
|
|
max_reward = 0
|
|
|
|
history_data = []
|
|
|
|
for i in range(100):
|
|
model.learn(total_timesteps=1500)
|
|
test_env = model.get_env()
|
|
test_obs = test_env.reset()
|
|
list_reward = []
|
|
# episolde_data = {'actions': [], 'observations': [], 'rewards': []}
|
|
|
|
for step in range(150):
|
|
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)
|
|
# episolde_data['actions'].append(test_action[0, :])
|
|
# episolde_data['observations'].append(np.array(test_obs))
|
|
# episolde_data['rewards'].append(test_rewards)
|
|
|
|
# history_data.append(episolde_data)
|
|
|
|
avg_reward = np.mean(list_reward[-130:])
|
|
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"))
|
|
|
|
# with open(os.path.join(parent_dir, "output", name + ".pkl"), 'wb') as f:
|
|
# pickle.dump(history_data, f) |