A REAL PROGRESS

CelerisLab/driver.py

@@ -90,7 +90,7 @@ class FlowField:
 
         # Initialize memory
         self.ddf = np.zeros(self.FIELD_SIZE * self.LATTICE, dtype=self.DATA_TYPE)
-        # self.velo = np.zeros(self.FIELD_SIZE * self.DIM, dtype=self.DATA_TYPE)
+        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)
 
@@ -278,5 +278,11 @@ class FlowField:
     def get_ddf(self):
         cuda.memcpy_dtoh(self.ddf, self.ddf_gpu)
 
+    def save_ddf(self):
+        self.ddf_save = self.ddf.copy()
+
+    def restore_ddf(self):
+        self.ddf = self.ddf_save.copy()
+
     def __del__(self):
         self.context.pop()

scripts/gym_env.py

@@ -120,7 +120,7 @@ class CustomEnv(gym.Env):
             # barrier.wait()
             self.fifo_states.append(self.flow_field.obs.copy()[2:14])
 
-        def calc_reward():
+        def proc_data():
             states = np.array(self.fifo_states)
             forces = states[-1, 6:12] / self.force_norm_fact
             cd = (forces[0] + forces[2] + forces[4]) / 3
@@ -179,7 +179,7 @@ class CustomEnv(gym.Env):
             result_queue.put((np.hstack([forces, sens]), reward))
 
         run_flow_field(action)
-        calc_reward()
+        proc_data()
         observation, reward = result_queue.get()
 
         truncated = bool(np.any(observation > 1) or np.any(observation < -1))