Major update: functional LBM and DRL

2024-08-18 19:09:09 +08:00
parent 89f2678579
commit c6f607dd8d
51 changed files with 3219 additions and 743 deletions
@@ -0,0 +1,57 @@
+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"))
@@ -0,0 +1,198 @@
+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 threading
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ProcessPoolExecutor
+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 = 800
+FIFO_LEN = 120
+CONV_LEN = 60
+MAX_STEPS = 640
+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, 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
+        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)
+        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()[2:8]
+            self.target_states = np.vstack((self.target_states, new_state))
+
+        self.flow_field.apply_ddf()
+        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)
+        self.flow_field.add_cylinder(center, L0 / 2)
+        center: Tuple[float, float, float] = (31.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))
+        self.flow_field.get_ddf()
+
+        for i in range(FIFO_LEN):
+            self.flow_field.run(SAMPLE_INTERVAL, np.zeros(7, dtype=DATA_TYPE))
+            self.fifo_states.append(self.flow_field.obs.copy()[2:14])
+
+        temp_states = np.array(self.fifo_states)
+        self.force_norm_fact = 6 * np.max(np.abs(temp_states[:, 6:12]))
+        for i in range(6):
+            self.sens_deviation[i] = np.mean(temp_states[:, i])
+            self.sens_norm_fact[i] = 5 * np.max(np.abs(temp_states[:, i] - self.sens_deviation[i]))
+            self.target_states[:, i] = (self.target_states[:, i] - 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(7, dtype=DATA_TYPE)
+                temp[4:7] = 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()[2:14])
+
+        def calc_reward():
+            states = np.array(self.fifo_states)
+            forces = states[-1, 6:12] / self.force_norm_fact
+            cd = (forces[0] + forces[2] + forces[4]) / 3
+            cl = (forces[1] + forces[3] + forces[5]) / 3
+            sens = (states[-1, 0:6] - self.sens_deviation) / self.sens_norm_fact
+
+            similarities = 0.0
+
+            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)
+
+            id_sens = 0
+            target_seq = self.target_states[:, id_sens]
+            state_seq = (states[:, id_sens] - self.sens_deviation[id_sens]) / self.sens_norm_fact[id_sens]
+            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]
+                state_seq = (states[:, i] - self.sens_deviation[i]) / self.sens_norm_fact[i]
+                similarities += calc_sim(target_seq, state_seq, lag) / 6
+
+            reward_cd = np.exp(-np.abs(cd * 80))
+            reward_cl = np.exp(-np.abs(cl * 20))
+            # reward_sim = np.exp(2 * (similarities - 1))
+            reward_sim = similarities
+            reward = np.minimum(0.3 * reward_cd + 0.3 * reward_cl + 0.4 * reward_sim, 1.0)
+            # barrier.wait()
+            result_queue.put((np.hstack([forces, sens]), reward))
+
+        run_flow_field(action)
+        calc_reward()
+        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"):
+        pass
+
+    def close(self):
+        self.flow_field.__del__()