74 lines
2.8 KiB
Python
74 lines
2.8 KiB
Python
import gymnasium as gym
|
|
import numpy as np
|
|
from gymnasium import spaces
|
|
import ctypes
|
|
from collections import deque
|
|
|
|
lbm = ctypes.cdll.LoadLibrary('./lbm_sens.so')
|
|
|
|
S_DIM, A_DIM = 6, 3
|
|
action_amp = 5
|
|
action_weight = 0.5
|
|
sample_interval = 200
|
|
max_steps = 320
|
|
|
|
class CustomEnv(gym.Env):
|
|
"""Custom Environment that follows gym interface."""
|
|
|
|
metadata = {"render_modes": ["human"], "render_fps": 1000/sample_interval}
|
|
|
|
def __init__(self, devicenum=0, Ccost=0.2):
|
|
super().__init__()
|
|
self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)
|
|
self.observation_space = spaces.Box(low=-5, high=5, shape=(6,), dtype=np.float32)
|
|
self.fifo_rewards = deque(maxlen=50)
|
|
lbm.SetDevice(devicenum)
|
|
lbm.InitAll()
|
|
lbm.CoreSolver.argtypes = (ctypes.c_int,ctypes.c_float,ctypes.c_float,ctypes.c_float,ctypes.c_float)
|
|
lbm.CoreSolver.restype = ctypes.POINTER(ctypes.c_float)
|
|
self.temps_init = lbm.CoreSolver(100*1000, 0.0, 0.0, 0.0, 0.0)
|
|
self.s = np.array([0.0] * S_DIM, dtype=np.float32)
|
|
for i in range(S_DIM):
|
|
self.s[i] = self.temps_init[i]
|
|
lbm.InitCPUMemory()
|
|
self.max_steps = max_steps
|
|
self.current_step = 0
|
|
self.Ccost = Ccost
|
|
|
|
def step(self, action):
|
|
assert self.action_space.contains(action), "%r (%s) invalid"%(action, type(action))
|
|
lbm.CoreSolver.argtypes = (ctypes.c_int,ctypes.c_float,ctypes.c_float,ctypes.c_float,ctypes.c_float)
|
|
lbm.CoreSolver.restype = ctypes.POINTER(ctypes.c_float)
|
|
action = action_amp * action
|
|
temps = lbm.CoreSolver(sample_interval, action_weight, action[0], action[1], action[2])
|
|
for i in range(S_DIM):
|
|
self.s[i] = temps[i]
|
|
observation = np.hstack(self.s)
|
|
cd = self.s[0]+self.s[2]+self.s[4]
|
|
cl = self.s[1]+self.s[3]+self.s[5]
|
|
reward = float((1-self.Ccost)*np.exp(-np.abs(cd)/3)+self.Ccost*np.exp(-np.abs(cl)/3))
|
|
self.fifo_rewards.append(reward)
|
|
terminated = bool(np.mean(self.fifo_rewards) > 0.9)
|
|
truncated = bool(np.any(self.s > 3) or np.any(self.s < -3))
|
|
self.current_step += 1
|
|
if self.current_step >= self.max_steps:
|
|
terminated = True
|
|
info = {}
|
|
return observation, reward, terminated, truncated, info
|
|
|
|
def reset(self, seed=None, Ccost=0.2):
|
|
lbm.ResetAll()
|
|
for i in range(S_DIM):
|
|
self.s[i] = self.temps_init[i]
|
|
observation = np.hstack(self.s)
|
|
info = {}
|
|
self.current_step = 0
|
|
self.Ccost = Ccost
|
|
return observation, info
|
|
|
|
def render(self, episode=0, numstep=0):
|
|
lbm.OutputFlow.argtypes = (ctypes.c_int, ctypes.c_int, ctypes.c_int)
|
|
lbm.OutputFlow(episode, numstep, sample_interval)
|
|
|
|
def close(self):
|
|
lbm.Finalize() |