Frank14f/Frank_LBM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

240923

Browse Source
This commit is contained in:
Frank14f
2024-09-23 22:36:33 +08:00
parent 729c5ee6ee
commit acf2b36b8c
29 changed files with 983696 additions and 318 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/__pycache__/gym_env_erase.cpython-310.pyc
BIN
View File
Binary file not shown.
scripts/__pycache__/gym_env_imit.cpython-310.pyc
BIN
View File
Binary file not shown.
scripts/d1a3o12_erase.py
+8 -8
View File
@@ -27,16 +27,16 @@ class Sin(Module):
if __name__ == '__main__':
vec_env = CustomEnv(device_id=1)
name = "d1a3o12_re100_erase_1"
name = "d1a3o12_re100_erase_c0"
model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_re100_erase"), env=vec_env, device=torch.device("cuda:1"))
# model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_re100_erase_b0"), 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)
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
scripts/d1a3o12_imit.py
+70
View File
@@ -0,0 +1,70 @@
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_imit 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 = "d1a3o12_re100_imit_a1"
model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_re100_imit_a0"), 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"),
# verbose=0)
writer = SummaryWriter(log_dir=os.path.join(parent_dir, "tensorboard", name))
max_reward = 0
history_data = []
for i in range(400):
model.learn(total_timesteps=360)
test_env = model.get_env()
test_obs = test_env.reset()
list_reward = []
episolde_data = {'actions': [], 'observations': [], 'rewards': []}
for step in range(360):
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[-180:])
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)
scripts/gym_env_erase.py
+9 -6
View File
@@ -28,7 +28,7 @@ config_field = utils.load_flow_field_config(
S_DIM, A_DIM = 12, 3
U0 = config_field.velocity
T0 = 1000
SAMPLE_INTERVAL = 800
SAMPLE_INTERVAL = 1200
FIFO_LEN = 120
CONV_LEN = 60
MAX_STEPS = 360
@@ -50,13 +50,14 @@ class CustomEnv(gym.Env):
low=-1, high=1, shape=(S_DIM,), dtype=DATA_TYPE
)
self.fifo_states = deque(maxlen=FIFO_LEN)
self.fifo_rewards = deque([0.1] * 10, maxlen=10)
self.target_states = np.empty((0, 6), dtype=DATA_TYPE)
self.force_norm_fact = 1.0
self.sens_norm_fact = np.ones(6, dtype=DATA_TYPE)
self.sens_deviation = np.zeros(6, dtype=DATA_TYPE)
self.flow_field = FlowField(config_field, config_cuda, device_id)
L0 = 20
L0 = 30
U0 = config_field.velocity
NX = self.flow_field.FIELD_SHAPE[0]
NY = self.flow_field.FIELD_SHAPE[1]
@@ -80,9 +81,9 @@ class CustomEnv(gym.Env):
self.flow_field.add_cylinder(center, L0)
center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
self.flow_field.add_cylinder(center, L0 / 2)
center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 + 0.75 * L0, 0)
center: Tuple[float, float, float] = ((30+1.3) * L0, (NY - 1) / 2 + 0.75 * L0, 0)
self.flow_field.add_cylinder(center, L0 / 2)
center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 - 0.75 * L0, 0)
center: Tuple[float, float, float] = ((30+1.3) * L0, (NY - 1) / 2 - 0.75 * L0, 0)
self.flow_field.add_cylinder(center, L0 / 2)
self.flow_field.run(int(4*NX/U0), np.zeros(7, dtype=DATA_TYPE))
@@ -113,7 +114,7 @@ class CustomEnv(gym.Env):
U0 = config_field.velocity
try:
temp = np.zeros(7, dtype=DATA_TYPE)
temp[4:7] = np.array((action*8+[0,-4,4])*U0, dtype=DATA_TYPE)
temp[4:7] = np.array((action*6+[0,-6,6])*U0, dtype=DATA_TYPE)
self.flow_field.run(SAMPLE_INTERVAL, temp)
finally:
self.flow_field.context.pop()
@@ -134,7 +135,9 @@ class CustomEnv(gym.Env):
reward_cd = np.exp(-np.abs(cd * 20))
reward_cl = np.exp(-np.abs(cl * 80))
reward_sim = np.exp(-np.abs(diff * 20))
reward = np.minimum(0.3 * reward_cd + 0.3 * reward_cl + 0.4 * reward_sim, 1.0)
reward = np.minimum(0.2 * reward_cd + 0.2 * reward_cl + 0.6 * reward_sim, 1.0)
self.fifo_rewards.append(reward)
reward = np.mean(self.fifo_rewards)
result_queue.put((np.hstack([forces[2:8], sens]), reward))
scripts/gym_env_imit.py
+237
View File
@@ -0,0 +1,237 @@
import gymnasium as gym
import numpy as np
from gymnasium import spaces
import ctypes
from collections import deque
from typing import Tuple
import sys
import os
import matplotlib.pyplot as plt
import queue
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
current_dir = os.path.dirname(os.path.abspath("__file__"))
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
sys.path.append(parent_dir)
from CelerisLab import FlowField
from CelerisLab import utils
config_cuda = utils.load_cuda_config(
os.path.join(parent_dir, "configs", "config_cuda.json")
)
config_field = utils.load_flow_field_config(
os.path.join(parent_dir, "configs", "config_flowfield.json")
)
S_DIM, A_DIM = 12, 3
U0 = config_field.velocity
T0 = 1000
SAMPLE_INTERVAL = 1200
FIFO_LEN = 120
CONV_LEN = 60
MAX_STEPS = 360
if config_field.data_type == "FP32":
DATA_TYPE = np.float32
else:
raise ValueError(f"Unsupported data type {config_field.data_type}.")
class CustomEnv(gym.Env):
"""Custom Environment that follows gym interface."""
metadata = {"render_modes": ["human"], "render_fps": T0 / SAMPLE_INTERVAL}
def __init__(self, device_id=0):
super().__init__()
self.action_space = spaces.Box(low=-1, high=1, shape=(A_DIM,), dtype=DATA_TYPE)
self.observation_space = spaces.Box(
low=-1, high=1, shape=(S_DIM,), dtype=DATA_TYPE
)
self.fifo_states = deque(maxlen=FIFO_LEN)
self.target_states = np.empty((0, 8), dtype=DATA_TYPE)
self.force_norm_fact = 1.0
self.sens_norm_fact = np.ones(6, dtype=DATA_TYPE)
self.sens_deviation = np.zeros(6, dtype=DATA_TYPE)
self.flow_field = FlowField(config_field, config_cuda, device_id)
L0 = 30
U0 = config_field.velocity
NX = self.flow_field.FIELD_SHAPE[0]
NY = self.flow_field.FIELD_SHAPE[1]
center: Tuple[float, float, float] = (10 * L0, (NY - 1) / 2, 0)
self.flow_field.add_cylinder(center, L0 / 2)
center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 + 2 * L0, 0)
self.flow_field.add_sensor(center, L0 / 4)
center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2, 0)
self.flow_field.add_sensor(center, L0 / 4)
center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 - 2 * L0, 0)
self.flow_field.add_sensor(center, L0 / 4)
self.flow_field.run(int(4*NX/U0), np.zeros(4, dtype=DATA_TYPE))
for i in range(FIFO_LEN):
self.flow_field.run(SAMPLE_INTERVAL, np.zeros(4, dtype=DATA_TYPE))
new_state = self.flow_field.obs.copy()
self.target_states = np.vstack((self.target_states, new_state))
self.flow_field.apply_ddf()
center: Tuple[float, float, float] = (10 * L0, (NY - 1) / 2, 0)
self.flow_field.add_cylinder(center, L0)
center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
self.flow_field.add_cylinder(center, L0 / 2)
center: Tuple[float, float, float] = ((30+1.3) * L0, (NY - 1) / 2 + 0.75 * L0, 0)
self.flow_field.add_cylinder(center, L0 / 2)
center: Tuple[float, float, float] = ((30+1.3) * L0, (NY - 1) / 2 - 0.75 * L0, 0)
self.flow_field.add_cylinder(center, L0 / 2)
self.flow_field.run(int(4*NX/U0), np.zeros(8, dtype=DATA_TYPE))
self.flow_field.get_ddf()
for i in range(FIFO_LEN):
self.flow_field.run(SAMPLE_INTERVAL, np.zeros(8, dtype=DATA_TYPE))
self.fifo_states.append(self.flow_field.obs.copy())
temp_states = np.array(self.fifo_states)
self.force_norm_fact = 6 * np.max(np.abs(temp_states[:, 8:16]))
for i in range(6):
self.sens_deviation[i] = np.mean(temp_states[:, i+2])
self.sens_norm_fact[i] = 5 * np.max(np.abs(temp_states[:, i+2] - self.sens_deviation[i]))
self.target_states[:, i+2] = (self.target_states[:, i+2] - self.sens_deviation[i]) / self.sens_norm_fact[i]
def step(self, action):
assert self.action_space.contains(action), "%r (%s) invalid" % (
action,
type(action),
)
# barrier = threading.Barrier(2)
result_queue = queue.Queue()
def run_flow_field(action):
self.flow_field.context.push()
U0 = config_field.velocity
try:
temp = np.zeros(8, dtype=DATA_TYPE)
temp[5:8] = np.array((action*8+[0,-4,4])*U0, dtype=DATA_TYPE)
self.flow_field.run(SAMPLE_INTERVAL, temp)
finally:
self.flow_field.context.pop()
# barrier.wait()
self.fifo_states.append(self.flow_field.obs.copy())
def proc_data():
states = np.array(self.fifo_states)
forces = states[-1, 8:16] / self.force_norm_fact
sens = (states[-1, 2:8] - self.sens_deviation) / self.sens_norm_fact
cd = forces[0] + forces[2] + forces[4] + forces[6]
cl = forces[1] + forces[3] + forces[5] + forces[7]
def calc_lag(target, state):
target_mean = np.mean(target)
state_mean = np.mean(state)
correlation = np.correlate(target - target_mean, state - state_mean, "full")
lags = np.arange(-len(target) + 1, len(target))
max_lag = lags[np.argmax(correlation)]
return max_lag
def calc_sim(target, state, lag):
target_mean = np.mean(target)
state_mean = np.mean(state)
target_std = np.std(target)
aligned_state = np.roll(state, lag)
if lag >= 0:
seq_target = target[-CONV_LEN:]-target_mean
seq_state = aligned_state[-CONV_LEN:]-state_mean
else:
seq_target = target[:CONV_LEN]-target_mean
seq_state = aligned_state[:CONV_LEN]-state_mean
seq_diff = seq_target - seq_state
sim_cor = 10*(np.corrcoef(seq_target, seq_state)[0, 1] - 1)
sim_div = -np.abs((target_mean - state_mean) / target_std * 0.75)
sim_amp = -np.abs(np.std(seq_diff) / target_std * 2)
return np.exp((sim_cor + sim_div + sim_amp) / 3)
similarities = 0.0
target_seq = self.target_states[:, 2]
state_seq = (states[:, 2] - self.sens_deviation[0]) / self.sens_norm_fact[0]
lag = calc_lag(target_seq, state_seq)
similarities += calc_sim(target_seq, state_seq, lag) / 6
for i in range(1, 6):
target_seq = self.target_states[:, i+2]
state_seq = (states[:, i+2] - self.sens_deviation[i]) / self.sens_norm_fact[i]
similarities += calc_sim(target_seq, state_seq, lag) / 6
reward_sim = similarities
target_seq = self.target_states[:, 0]
state_seq = states[:, 8] + states[:, 10] + states[:, 12] + states[:, 14]
ave_drag = np.average(state_seq)
lag = calc_lag(target_seq, state_seq)
similarities += calc_sim(target_seq, state_seq, lag) / 2
target_seq = self.target_states[:, 1]
state_seq = states[:, 9] + states[:, 11] + states[:, 13] + states[:, 15]
similarities += calc_sim(target_seq, state_seq, lag) / 2
reward_force = similarities
reward_cd = np.exp(-np.abs((cd-ave_drag) * 2))
reward_cl = np.exp(-np.abs(cl * 8))
reward = np.minimum(0.0 * reward_cd + 0.0 * reward_cl + 0.4 * reward_force + 0.6 * reward_sim, 1.0)
# barrier.wait()
result_queue.put((np.hstack([forces[2:8], sens]), reward))
run_flow_field(action)
proc_data()
observation, reward = result_queue.get()
truncated = bool(np.any(observation > 1) or np.any(observation < -1))
observation = np.clip(observation, -1, 1)
# truncated = False
return observation, float(reward), False, truncated, {}
def reset(self, seed=None):
self.flow_field.apply_ddf()
return np.zeros(S_DIM, dtype=np.float32), {}
def render(self, mode="human"):
NX = self.flow_field.FIELD_SHAPE[0]
NY = self.flow_field.FIELD_SHAPE[1]
self.flow_field.get_ddf()
ddf_plot = self.flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)
ux = (ddf_plot[:, :, 1] + ddf_plot[:, :, 5] + ddf_plot[:, :, 8] - ddf_plot[:, :, 3] - ddf_plot[:, :, 6] - ddf_plot[:, :, 7]) / U0
uy = (ddf_plot[:, :, 2] + ddf_plot[:, :, 5] + ddf_plot[:, :, 6] - ddf_plot[:, :, 4] - ddf_plot[:, :, 7] - ddf_plot[:, :, 8]) / U0
speed = np.sqrt(ux**2 + uy**2)
plt.figure(figsize=(10, 5))
plt.imshow(speed.T, origin='lower', cmap='viridis', extent=[0, NX, 0, NY])
plt.colorbar(label='Speed')
plt.title('Scalar Velocity Field')
plt.xlabel('X')
plt.ylabel('Y')
plt.tight_layout()
plt.show()
def save_field(self, filename):
NX = self.flow_field.FIELD_SHAPE[0]
NY = self.flow_field.FIELD_SHAPE[1]
self.flow_field.get_ddf()
ddf_plot = self.flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)
flag_plot = self.flow_field.flag.copy().reshape((NY, NX)).transpose(1, 0)
ux = (ddf_plot[:, :, 1] + ddf_plot[:, :, 5] + ddf_plot[:, :, 8] - ddf_plot[:, :, 3] - ddf_plot[:, :, 6] - ddf_plot[:, :, 7]) / U0
uy = (ddf_plot[:, :, 2] + ddf_plot[:, :, 5] + ddf_plot[:, :, 6] - ddf_plot[:, :, 4] - ddf_plot[:, :, 7] - ddf_plot[:, :, 8]) / U0
with open(os.path.join(parent_dir, "output", filename), "w") as f:
f.write("Title= \"LBM 2D\"\r\n")
f.write("VARIABLES= \"X\",\"Y\",\"flag\",\"U\",\"V\",\r\n")
f.write(f"ZONE T= \"BOX\",I= {NX},J= {NY},F=POINT\r\n")
for j in range(NY):
for i in range(NX):
f.write(f"{i},{j},{flag_plot[i, j]},{ux[i, j]},{uy[i, j]}\r\n")
def close(self):
self.flow_field.__del__()
scripts/nohup1.out
+16
View File
@@ -0,0 +1,16 @@
Traceback (most recent call last):
File "/home/frank14f/Frank_LBM/scripts/d1a3o12_imit.py", line 55, in <module>
test_obs, test_rewards, test_dones, info = test_env.step(test_action)
File "/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/base_vec_env.py", line 206, in step
return self.step_wait()
File "/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/dummy_vec_env.py", line 58, in step_wait
obs, self.buf_rews[env_idx], terminated, truncated, self.buf_infos[env_idx] = self.envs[env_idx].step(
File "/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/monitor.py", line 94, in step
observation, reward, terminated, truncated, info = self.env.step(action)
File "/home/frank14f/Frank_LBM/scripts/gym_env_imit.py", line 190, in step
run_flow_field(action)
File "/home/frank14f/Frank_LBM/scripts/gym_env_imit.py", line 117, in run_flow_field
self.flow_field.run(SAMPLE_INTERVAL, temp)
File "/home/frank14f/Frank_LBM/CelerisLab/driver.py", line 254, in run
cuda.memcpy_htod_async(self.action_gpu, action_pinned, stream)
KeyboardInterrupt