new branch test

.gitignore

@@ -1,2 +1,3 @@
 tensorboard/*
 models/*
+output/*

CelerisLab/driver.py

@@ -2,7 +2,7 @@
 
 import pycuda.driver as cuda
 import numpy as np
-from typing import List, Tuple, Union
+from typing import List, Tuple, Union, Optional
 
 from . import utils
 from . import preprocess as preproc
@@ -93,13 +93,13 @@ class FlowField:
         self.ddf_save = np.zeros(self.FIELD_SIZE * self.LATTICE, dtype=self.DATA_TYPE)
         self.flag = np.ones(self.FIELD_SIZE, dtype=np.uint8)
         self.indx = np.zeros(self.FIELD_SIZE, dtype=np.int32)
-
         self.delta_curve = np.zeros(0, dtype=self.DATA_TYPE)
 
         self.ddf_gpu = cuda.mem_alloc(self.ddf.nbytes)
         self.temp_gpu = cuda.mem_alloc(self.ddf.nbytes)
         self.flag_gpu = cuda.mem_alloc(self.flag.nbytes)
         self.indx_gpu = cuda.mem_alloc(self.indx.nbytes)
+        self.delta_gpu = cuda.mem_alloc(1)
 
         self.objects = {}
         self.action = np.zeros(0, dtype=self.DATA_TYPE)
@@ -118,7 +118,7 @@ class FlowField:
         cuda.memcpy_dtoh(self.flag, self.flag_gpu)
         cuda.memcpy_dtoh(self.ddf, self.ddf_gpu)
 
-    def add_cylinder(self, center: Tuple[float, float, float], radius: float):
+    def add_cylinder(self, center: Tuple[float, float, float], radius: float, id_obj: Optional[int] = None):
         x_c, y_c, z_c = center
 
         if (
@@ -130,10 +130,13 @@ class FlowField:
             raise ValueError("Cylinder is out of bounds.")
 
         index = self.delta_curve.size if self.delta_curve.size > 0 else 0
+
         if self.DATA_TYPE == np.float32:
             id_object = np.int32(len(self.objects))
+            # max_id = max(self.objects.keys())
         else:
             raise ValueError(f"Unsupported data type {self.DATA_TYPE}.")
+        
         for x in range(int(x_c - radius) - 1, int(x_c + radius) + 1):
             for y in range(int(y_c - radius) - 1, int(y_c + radius) + 1):
                 if (x - x_c) ** 2 + (y - y_c) ** 2 < radius**2:

scripts/d0a3o12.py

@@ -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_1 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 = "d0a3o12_b0"
+
+    # 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(240):
+        model.learn(total_timesteps=240)
+        test_env = model.get_env()
+        test_obs = test_env.reset()
+        list_reward = []
+        for step in range(240):
+            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[-120:])
+        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"))

scripts/d1a3o12.py

@@ -12,6 +12,7 @@ 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))
@@ -25,33 +26,45 @@ class Sin(Module):
 
 if __name__ == '__main__':
     
-    vec_env = CustomEnv(device_id=1)
+    vec_env = CustomEnv(device_id=3)
-    name = "d1a3o12_c1"
+    name = "d1a3o12_re100_new_reward"
 
-    model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_a0"), env=vec_env, device=torch.device("cuda:1"))
+    # model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_c0"), env=vec_env, device=torch.device("cuda:1"))
 
-    # model = PPO(
+    model = PPO(
-    #     "MlpPolicy", 
+        "MlpPolicy", 
-    #     policy_kwargs=dict(activation_fn=Sin),
+        policy_kwargs=dict(activation_fn=Sin),
-    #     env=vec_env, 
+        env=vec_env, 
-    #     device=torch.device("cuda:1"),
+        device=torch.device("cuda:3"),
-    #     verbose=0)
+        verbose=0)
 
     writer = SummaryWriter(log_dir=os.path.join(parent_dir, "tensorboard", name))
     max_reward = 0
 
-    for i in range(100):
+    history_data = []
-        model.learn(total_timesteps=480)
+
+    for i in range(400):
+        model.learn(total_timesteps=360)
         test_env = model.get_env()
         test_obs = test_env.reset()
         list_reward = []
-        for step in range(480):
+        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)
 
-        avg_reward = np.mean(list_reward[-240:])
+        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/d1a3o12_erase.py

@@ -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_erase 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=1)
+    name = "d1a3o12_re100_erase_1"
+
+    model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_re100_erase"), 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
+
+    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.py

@@ -6,9 +6,7 @@ from collections import deque
 from typing import Tuple
 import sys
 import os
-import threading
+import matplotlib.pyplot as plt
-from concurrent.futures import ThreadPoolExecutor
-from concurrent.futures import ProcessPoolExecutor
 import queue
 
 os.environ["OMP_NUM_THREADS"] = "1"
@@ -33,7 +31,7 @@ T0 = 1000
 SAMPLE_INTERVAL = 800
 FIFO_LEN = 120
 CONV_LEN = 60
-MAX_STEPS = 640
+MAX_STEPS = 360
 if config_field.data_type == "FP32":
     DATA_TYPE = np.float32
 else:
@@ -170,8 +168,8 @@ class CustomEnv(gym.Env):
                 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_cd = np.exp(-np.abs(cd * 20))
-            reward_cl = np.exp(-np.abs(cl * 20))
+            reward_cl = np.exp(-np.abs(cl * 80))
             # 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)
@@ -192,7 +190,37 @@ class CustomEnv(gym.Env):
         return np.zeros(S_DIM, dtype=np.float32), {}
     
     def render(self, mode="human"):
-        pass
+        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/gym_env_1.py

@@ -0,0 +1,188 @@
+import gymnasium as gym
+import numpy as np
+from gymnasium import spaces
+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 = 800
+FIFO_LEN = 120
+CONV_LEN = 60
+MAX_STEPS = 240
+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.force_norm_fact = 1.0
+        self.sens_norm_fact = 1.0
+        self.sens_deviation = np.zeros(6, dtype=DATA_TYPE)
+        self.current_step = 0
+
+        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] = (30 * L0, (NY - 1) / 2 + 2 * L0, 0)
+        self.flow_field.add_sensor(center, L0 / 4)
+        center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
+        self.flow_field.add_sensor(center, L0 / 4)
+        center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2 - 2 * L0, 0)
+        self.flow_field.add_sensor(center, L0 / 4)
+        center: Tuple[float, float, float] = (20 * L0, (NY - 1) / 2, 0)
+        self.flow_field.add_cylinder(center, L0 / 2)
+        center: Tuple[float, float, float] = (21.3 * L0, (NY - 1) / 2 + 0.75 * L0, 0)
+        self.flow_field.add_cylinder(center, L0 / 2)
+        center: Tuple[float, float, float] = (21.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(6, dtype=DATA_TYPE))
+        self.flow_field.get_ddf()
+        self.flow_field.save_ddf()
+
+        # for i in range(FIFO_LEN):
+        #     self.flow_field.run(SAMPLE_INTERVAL, np.zeros(6, dtype=DATA_TYPE))
+        #     self.fifo_states.append(self.flow_field.obs.copy())
+
+        # temp_states = np.array(self.fifo_states)
+        # self.force_norm_fact = 3 * np.max(np.abs(temp_states[:, 6:12]))
+        # for i in range(3):
+        #     self.sens_deviation[2*i] = np.mean(temp_states[:, 2*i])
+        #     # self.sens_deviation[2*i] = 0.0
+        # self.sens_norm_fact = 6 * np.max(np.abs(temp_states[:, 1]))
+            
+
+    def step(self, action):
+        assert self.action_space.contains(action), "%r (%s) invalid" % (
+            action,
+            type(action),
+        )
+
+        result_queue = queue.Queue()
+
+        def run_flow_field(action):
+            self.flow_field.context.push()
+            U0 = config_field.velocity
+            try:
+                temp = np.zeros(6, dtype=DATA_TYPE)
+                temp[3:6] = np.array((action*8)*U0, dtype=DATA_TYPE)
+                self.flow_field.run(SAMPLE_INTERVAL, temp)
+            finally:
+                self.fifo_states.append(self.flow_field.obs.copy())
+                self.flow_field.context.pop()
+
+        def proc_data():
+            states = np.array(self.fifo_states)
+            forces = states[-1, 6:12] * 50
+            cd = (forces[0] + forces[2] + forces[4]) / 3
+            cl = (forces[1] + forces[3] + forces[5]) / 3
+            sens = states[-1, 0:6] / 2
+
+            def theo_velo(y):
+                NY = self.flow_field.FIELD_SHAPE[1]
+                U0 = config_field.velocity
+                yy = (y - 0.5 * (NY - 1)) / (NY - 2.0)
+                u = U0 * 1.5 * (1 - 4 * yy * yy)
+                return u
+            
+            similarities = 0.0
+            NY = self.flow_field.FIELD_SHAPE[1]
+            L0 = 20
+            sens_pos = np.array([(NY - 1) / 2 + 2 * L0, (NY - 1) / 2, (NY - 1) / 2 - 2 * L0])
+            for i in range(3):
+                u = theo_velo(sens_pos[i])
+                similarities += np.exp(-4*np.abs(states[-1, 2*i] - u))/6 + np.exp(-8*np.abs(states[-1, 2*i+1] - 0))/6
+
+            reward_cd = np.exp(-np.abs(cd * 10))
+            reward_cl = np.exp(-np.abs(cl * 10))
+            # reward_sim = np.exp(2 * (similarities - 1))
+            reward_sim = similarities
+            reward = np.minimum(0.4 * reward_cd + 0.6 * reward_cl + 0.0 * reward_sim, 1.0)
+            # barrier.wait()
+            result_queue.put((np.hstack([forces, 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)
+        terminated = self.current_step >= MAX_STEPS
+        self.current_step += 1
+        return observation, float(reward), terminated, truncated, {}
+    
+    def reset(self, seed=None):
+        self.flow_field.apply_ddf()
+        self.current_step = 0
+        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/gym_env_erase.py

@@ -0,0 +1,188 @@
+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 = 800
+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, 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] = (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(1*NX/U0), np.zeros(3, dtype=DATA_TYPE))
+
+        for i in range(FIFO_LEN):
+            self.flow_field.run(SAMPLE_INTERVAL, np.zeros(3, dtype=DATA_TYPE))
+            new_state = self.flow_field.obs.copy()
+            self.target_states = np.vstack((self.target_states, new_state))
+
+        self.target_states = np.mean(self.target_states, axis=0)
+
+        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] = (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())
+
+        temp_states = np.array(self.fifo_states)
+        self.force_norm_fact = 6 * np.max(np.abs(temp_states[:, 8:14]))
+        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),
+        )
+
+        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()
+            self.fifo_states.append(self.flow_field.obs.copy())
+
+        def proc_data():
+            states = np.array(self.fifo_states)
+            forces = states[-1, 6:14] / self.force_norm_fact
+            cd = (forces[0] + forces[2] + forces[4] + forces[6]) / 6
+            cl = (forces[1] + forces[3] + forces[5] + forces[7]) / 6
+            sens = (states[-1, 0:6] - self.sens_deviation) / self.sens_norm_fact
+
+            diff = 0
+            for i in range(1, 6):
+                target = self.target_states[i]
+                diff += np.abs(sens[i] - target) / 6
+
+            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)
+
+            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__()