Add new files and make code improvements
This commit is contained in:
Frank14f
Binary file not shown.
@@ -0,0 +1,74 @@
|
||||
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()
|
||||
@@ -0,0 +1,148 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Using cuda device\n",
|
||||
"Logging to ./tensorboard/PPO_1\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Process ForkServerProcess-2:\n",
|
||||
"Process ForkServerProcess-1:\n",
|
||||
"Traceback (most recent call last):\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
|
||||
" self.run()\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
|
||||
" self._target(*self._args, **self._kwargs)\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
|
||||
" observation, reward, terminated, truncated, info = env.step(data)\n",
|
||||
"ValueError: not enough values to unpack (expected 5, got 4)\n",
|
||||
"Traceback (most recent call last):\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
|
||||
" self.run()\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
|
||||
" self._target(*self._args, **self._kwargs)\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
|
||||
" observation, reward, terminated, truncated, info = env.step(data)\n",
|
||||
"ValueError: not enough values to unpack (expected 5, got 4)\n",
|
||||
"Process ForkServerProcess-4:\n",
|
||||
"Traceback (most recent call last):\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
|
||||
" self.run()\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
|
||||
" self._target(*self._args, **self._kwargs)\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
|
||||
" observation, reward, terminated, truncated, info = env.step(data)\n",
|
||||
"ValueError: not enough values to unpack (expected 5, got 4)\n",
|
||||
"Process ForkServerProcess-3:\n",
|
||||
"Traceback (most recent call last):\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
|
||||
" self.run()\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
|
||||
" self._target(*self._args, **self._kwargs)\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
|
||||
" observation, reward, terminated, truncated, info = env.step(data)\n",
|
||||
"ValueError: not enough values to unpack (expected 5, got 4)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"ename": "EOFError",
|
||||
"evalue": "",
|
||||
"output_type": "error",
|
||||
"traceback": [
|
||||
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
|
||||
"\u001b[0;31mEOFError\u001b[0m Traceback (most recent call last)",
|
||||
"Cell \u001b[0;32mIn[1], line 24\u001b[0m\n\u001b[1;32m 16\u001b[0m vec_env \u001b[38;5;241m=\u001b[39m SubprocVecEnv(env_fns)\n\u001b[1;32m 18\u001b[0m model \u001b[38;5;241m=\u001b[39m PPO(\n\u001b[1;32m 19\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMlpPolicy\u001b[39m\u001b[38;5;124m\"\u001b[39m, \n\u001b[1;32m 20\u001b[0m env\u001b[38;5;241m=\u001b[39mvec_env, \n\u001b[1;32m 21\u001b[0m n_steps\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m64\u001b[39m,\n\u001b[1;32m 22\u001b[0m tensorboard_log\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m./tensorboard/\u001b[39m\u001b[38;5;124m\"\u001b[39m, \n\u001b[1;32m 23\u001b[0m verbose\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m---> 24\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtotal_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m128\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m1000\u001b[39;49m\u001b[43m)\u001b[49m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/ppo/ppo.py:315\u001b[0m, in \u001b[0;36mPPO.learn\u001b[0;34m(self, total_timesteps, callback, log_interval, tb_log_name, reset_num_timesteps, progress_bar)\u001b[0m\n\u001b[1;32m 306\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mlearn\u001b[39m(\n\u001b[1;32m 307\u001b[0m \u001b[38;5;28mself\u001b[39m: SelfPPO,\n\u001b[1;32m 308\u001b[0m total_timesteps: \u001b[38;5;28mint\u001b[39m,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 313\u001b[0m progress_bar: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m 314\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m SelfPPO:\n\u001b[0;32m--> 315\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 316\u001b[0m \u001b[43m \u001b[49m\u001b[43mtotal_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtotal_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 317\u001b[0m \u001b[43m \u001b[49m\u001b[43mcallback\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcallback\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 318\u001b[0m \u001b[43m \u001b[49m\u001b[43mlog_interval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlog_interval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 319\u001b[0m \u001b[43m \u001b[49m\u001b[43mtb_log_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtb_log_name\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 320\u001b[0m \u001b[43m \u001b[49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 321\u001b[0m \u001b[43m \u001b[49m\u001b[43mprogress_bar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprogress_bar\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 322\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/on_policy_algorithm.py:277\u001b[0m, in \u001b[0;36mOnPolicyAlgorithm.learn\u001b[0;34m(self, total_timesteps, callback, log_interval, tb_log_name, reset_num_timesteps, progress_bar)\u001b[0m\n\u001b[1;32m 274\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39menv \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 276\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_timesteps \u001b[38;5;241m<\u001b[39m total_timesteps:\n\u001b[0;32m--> 277\u001b[0m continue_training \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcollect_rollouts\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menv\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallback\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrollout_buffer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_rollout_steps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mn_steps\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 279\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m continue_training:\n\u001b[1;32m 280\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/on_policy_algorithm.py:194\u001b[0m, in \u001b[0;36mOnPolicyAlgorithm.collect_rollouts\u001b[0;34m(self, env, callback, rollout_buffer, n_rollout_steps)\u001b[0m\n\u001b[1;32m 189\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 190\u001b[0m \u001b[38;5;66;03m# Otherwise, clip the actions to avoid out of bound error\u001b[39;00m\n\u001b[1;32m 191\u001b[0m \u001b[38;5;66;03m# as we are sampling from an unbounded Gaussian distribution\u001b[39;00m\n\u001b[1;32m 192\u001b[0m clipped_actions \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mclip(actions, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maction_space\u001b[38;5;241m.\u001b[39mlow, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maction_space\u001b[38;5;241m.\u001b[39mhigh)\n\u001b[0;32m--> 194\u001b[0m new_obs, rewards, dones, infos \u001b[38;5;241m=\u001b[39m \u001b[43menv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[43mclipped_actions\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 196\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_timesteps \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m env\u001b[38;5;241m.\u001b[39mnum_envs\n\u001b[1;32m 198\u001b[0m \u001b[38;5;66;03m# Give access to local variables\u001b[39;00m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/base_vec_env.py:206\u001b[0m, in \u001b[0;36mVecEnv.step\u001b[0;34m(self, actions)\u001b[0m\n\u001b[1;32m 199\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 200\u001b[0m \u001b[38;5;124;03mStep the environments with the given action\u001b[39;00m\n\u001b[1;32m 201\u001b[0m \n\u001b[1;32m 202\u001b[0m \u001b[38;5;124;03m:param actions: the action\u001b[39;00m\n\u001b[1;32m 203\u001b[0m \u001b[38;5;124;03m:return: observation, reward, done, information\u001b[39;00m\n\u001b[1;32m 204\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 205\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstep_async(actions)\n\u001b[0;32m--> 206\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep_wait\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py:129\u001b[0m, in \u001b[0;36mSubprocVecEnv.step_wait\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstep_wait\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m VecEnvStepReturn:\n\u001b[0;32m--> 129\u001b[0m results \u001b[38;5;241m=\u001b[39m [remote\u001b[38;5;241m.\u001b[39mrecv() \u001b[38;5;28;01mfor\u001b[39;00m remote \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mremotes]\n\u001b[1;32m 130\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mwaiting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m 131\u001b[0m obs, rews, dones, infos, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreset_infos \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;241m*\u001b[39mresults) \u001b[38;5;66;03m# type: ignore[assignment]\u001b[39;00m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py:129\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstep_wait\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m VecEnvStepReturn:\n\u001b[0;32m--> 129\u001b[0m results \u001b[38;5;241m=\u001b[39m [\u001b[43mremote\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrecv\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m remote \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mremotes]\n\u001b[1;32m 130\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mwaiting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m 131\u001b[0m obs, rews, dones, infos, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreset_infos \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;241m*\u001b[39mresults) \u001b[38;5;66;03m# type: ignore[assignment]\u001b[39;00m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/connection.py:250\u001b[0m, in \u001b[0;36m_ConnectionBase.recv\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 248\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_closed()\n\u001b[1;32m 249\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_readable()\n\u001b[0;32m--> 250\u001b[0m buf \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_recv_bytes\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 251\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _ForkingPickler\u001b[38;5;241m.\u001b[39mloads(buf\u001b[38;5;241m.\u001b[39mgetbuffer())\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/connection.py:414\u001b[0m, in \u001b[0;36mConnection._recv_bytes\u001b[0;34m(self, maxsize)\u001b[0m\n\u001b[1;32m 413\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_recv_bytes\u001b[39m(\u001b[38;5;28mself\u001b[39m, maxsize\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m--> 414\u001b[0m buf \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_recv\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 415\u001b[0m size, \u001b[38;5;241m=\u001b[39m struct\u001b[38;5;241m.\u001b[39munpack(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m!i\u001b[39m\u001b[38;5;124m\"\u001b[39m, buf\u001b[38;5;241m.\u001b[39mgetvalue())\n\u001b[1;32m 416\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m size \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m:\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/connection.py:383\u001b[0m, in \u001b[0;36mConnection._recv\u001b[0;34m(self, size, read)\u001b[0m\n\u001b[1;32m 381\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m n \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m 382\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m remaining \u001b[38;5;241m==\u001b[39m size:\n\u001b[0;32m--> 383\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mEOFError\u001b[39;00m\n\u001b[1;32m 384\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 385\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mOSError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mgot end of file during message\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
|
||||
"\u001b[0;31mEOFError\u001b[0m: "
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"os.environ['MKL_THREADING_LAYER'] = 'GNU'\n",
|
||||
"import numpy as np\n",
|
||||
"import gymnasium as gym\n",
|
||||
"from env_pinball import CustomEnv\n",
|
||||
"from stable_baselines3 import PPO\n",
|
||||
"from stable_baselines3.common.vec_env import SubprocVecEnv\n",
|
||||
"\n",
|
||||
"def make_env(gpu_id):\n",
|
||||
" def _init():\n",
|
||||
" os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id)\n",
|
||||
" return CustomEnv(devicenum=gpu_id)\n",
|
||||
" return _init\n",
|
||||
"\n",
|
||||
"env_fns = [make_env(i) for i in range(4)]\n",
|
||||
"vec_env = SubprocVecEnv(env_fns)\n",
|
||||
"\n",
|
||||
"model = PPO(\n",
|
||||
" \"MlpPolicy\", \n",
|
||||
" env=vec_env, \n",
|
||||
" n_steps=64,\n",
|
||||
" tensorboard_log=\"./tensorboard/\", \n",
|
||||
" verbose=1)\n",
|
||||
"model.learn(total_timesteps=64*1000)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"vec_env = model.get_env()\n",
|
||||
"obs = vec_env.reset()\n",
|
||||
"\n",
|
||||
"n_steps = 0\n",
|
||||
"list_reward = {}\n",
|
||||
"terminated = False\n",
|
||||
"truncated = False\n",
|
||||
"while n_steps < 500 and not terminated and not truncated:\n",
|
||||
" n_steps += 1\n",
|
||||
" action, _states = model.predict(observation=obs)\n",
|
||||
" obs, rewards, dones, info = vec_env.step(action)\n",
|
||||
" list_reward[n_steps] = rewards"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "pycuda_3_10",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
Executable
BIN
Binary file not shown.
+216
@@ -0,0 +1,216 @@
|
||||
# %%
|
||||
#!/usr/bin/env python3
|
||||
from env_pinball import CustomEnv
|
||||
import os
|
||||
import pickle
|
||||
import random
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributions.normal import Normal
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
env = CustomEnv(devicenum=3)
|
||||
writer = SummaryWriter(log_dir='./tensorboard/DRL')
|
||||
|
||||
# %%
|
||||
class Policy_Network(nn.Module):
|
||||
"""Parametrized Policy Network."""
|
||||
|
||||
def __init__(self, obs_space_dims: int, action_space_dims: int):
|
||||
"""Initializes a neural network that estimates the mean and standard deviation
|
||||
of a normal distribution from which an action is sampled from.
|
||||
|
||||
Args:
|
||||
obs_space_dims: Dimension of the observation space
|
||||
action_space_dims: Dimension of the action space
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
hidden_space1 = 256 # Nothing special with 16, feel free to change
|
||||
hidden_space2 = 256 # Nothing special with 32, feel free to change
|
||||
|
||||
# Shared Network
|
||||
self.shared_net = nn.Sequential(
|
||||
nn.Linear(obs_space_dims, hidden_space1),
|
||||
nn.Tanh(),
|
||||
nn.Linear(hidden_space1, hidden_space2),
|
||||
nn.Tanh(),
|
||||
)
|
||||
|
||||
# Policy Mean specific Linear Layer
|
||||
self.policy_mean_net = nn.Sequential(
|
||||
nn.Linear(hidden_space2, action_space_dims)
|
||||
)
|
||||
|
||||
# Policy Std Dev specific Linear Layer
|
||||
self.policy_stddev_net = nn.Sequential(
|
||||
nn.Linear(hidden_space2, action_space_dims)
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Conditioned on the observation, returns the mean and standard deviation
|
||||
of a normal distribution from which an action is sampled from.
|
||||
|
||||
Args:
|
||||
x: Observation from the environment
|
||||
|
||||
Returns:
|
||||
action_means: predicted mean of the normal distribution
|
||||
action_stddevs: predicted standard deviation of the normal distribution
|
||||
"""
|
||||
shared_features = self.shared_net(x.float())
|
||||
|
||||
action_means = self.policy_mean_net(shared_features)
|
||||
action_means = torch.tanh(action_means)
|
||||
|
||||
action_stddevs = torch.log(
|
||||
1 + torch.exp(self.policy_stddev_net(shared_features))
|
||||
)
|
||||
|
||||
return action_means, action_stddevs
|
||||
|
||||
# %%
|
||||
class REINFORCE:
|
||||
"""REINFORCE algorithm."""
|
||||
|
||||
def __init__(self, obs_space_dims: int, action_space_dims: int):
|
||||
"""Initializes an agent that learns a policy via REINFORCE algorithm [1]
|
||||
to solve the task at hand (Inverted Pendulum v4).
|
||||
|
||||
Args:
|
||||
obs_space_dims: Dimension of the observation space
|
||||
action_space_dims: Dimension of the action space
|
||||
"""
|
||||
|
||||
# Hyperparameters
|
||||
self.learning_rate = 1e-4 # Learning rate for policy optimization
|
||||
self.gamma = 0.99 # Discount factor
|
||||
self.eps = 1e-6 # small number for mathematical stability
|
||||
|
||||
self.probs = [] # Stores probability values of the sampled action
|
||||
self.rewards = [] # Stores the corresponding rewards
|
||||
|
||||
self.net = Policy_Network(obs_space_dims, action_space_dims)
|
||||
self.optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.learning_rate)
|
||||
|
||||
def sample_action(self, state: np.ndarray) -> float:
|
||||
"""Returns an action, conditioned on the policy and observation.
|
||||
|
||||
Args:
|
||||
state: Observation from the environment
|
||||
|
||||
Returns:
|
||||
action: Action to be performed
|
||||
"""
|
||||
state = torch.tensor(np.array([state]))
|
||||
action_means, action_stddevs = self.net(state)
|
||||
|
||||
# create a normal distribution from the predicted
|
||||
# mean and standard deviation and sample an action
|
||||
distrib = Normal(action_means[0] + self.eps, action_stddevs[0] + self.eps)
|
||||
action = distrib.sample()
|
||||
prob = distrib.log_prob(action)
|
||||
|
||||
action = torch.tanh(action)
|
||||
action = action.numpy()
|
||||
|
||||
self.probs.append(prob)
|
||||
|
||||
return action
|
||||
|
||||
def update(self):
|
||||
"""Updates the policy network's weights."""
|
||||
running_g = 0
|
||||
gs = []
|
||||
|
||||
# Discounted return (backwards) - [::-1] will return an array in reverse
|
||||
for R in self.rewards[::-1]:
|
||||
running_g = R + self.gamma * running_g
|
||||
gs.insert(0, running_g)
|
||||
|
||||
deltas = torch.tensor(gs)
|
||||
|
||||
loss = 0
|
||||
# minimize -1 * prob * reward obtained
|
||||
for log_prob, delta in zip(self.probs, deltas):
|
||||
loss += log_prob.mean() * delta * (-1)
|
||||
|
||||
# Update the policy network
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
self.optimizer.step()
|
||||
|
||||
# Empty / zero out all episode-centric/related variables
|
||||
self.probs = []
|
||||
self.rewards = []
|
||||
|
||||
# %%
|
||||
total_num_episodes = int(5e3) # Total number of episodes
|
||||
obs_space_dims = 6
|
||||
action_space_dims = 3
|
||||
rewards_over_seeds = []
|
||||
MAX_REWARD = 0
|
||||
|
||||
# Check if there is a saved state
|
||||
if os.path.exists('saved_state.pkl'):
|
||||
with open('saved_state.pkl', 'rb') as f:
|
||||
i_seed, episode, agent, reward_over_episodes, rewards_over_seeds, MAX_REWARD = pickle.load(f)
|
||||
os.remove('saved_state.pkl') # Remove the saved state
|
||||
else:
|
||||
i_seed = 0
|
||||
episode = 0
|
||||
agent = None
|
||||
reward_over_episodes = None
|
||||
|
||||
for seed in [1][i_seed:]: # Fibonacci seeds
|
||||
# set seed
|
||||
torch.manual_seed(seed)
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
|
||||
# Reinitialize agent every seed
|
||||
if agent is None or reward_over_episodes is None:
|
||||
agent = REINFORCE(obs_space_dims, action_space_dims)
|
||||
reward_over_episodes = []
|
||||
|
||||
while episode < total_num_episodes+1:
|
||||
obs, info = env.reset(Ccost=0.2+episode/total_num_episodes*0.6)
|
||||
steps = 0
|
||||
done = False
|
||||
terminated = False
|
||||
truncated = False
|
||||
reward_over_steps = []
|
||||
while not done:
|
||||
action = agent.sample_action(obs)
|
||||
obs, reward, terminated, truncated, info = env.step(action)
|
||||
agent.rewards.append(reward)
|
||||
reward_over_steps.append(reward)
|
||||
steps += 1
|
||||
done = terminated or truncated
|
||||
|
||||
avg_reward = np.mean(reward_over_steps[-64:])
|
||||
reward_over_episodes.append(np.array([avg_reward], dtype=np.float32))
|
||||
agent.update()
|
||||
|
||||
if episode % 10 == 0:
|
||||
# print("Episode:", episode, "Average Reward:", int(avg_reward))
|
||||
writer.add_scalar('Average Reward', int(avg_reward), episode)
|
||||
|
||||
if avg_reward > MAX_REWARD:
|
||||
MAX_REWARD = avg_reward
|
||||
with open('saved_model_'+str(seed)+'.pkl', 'wb') as f:
|
||||
pickle.dump((episode + 1, agent, reward_over_episodes, MAX_REWARD), f)
|
||||
|
||||
# Save the current state at the end of each episode
|
||||
with open('saved_state.pkl', 'wb') as f:
|
||||
pickle.dump((i_seed, episode + 1, agent, reward_over_episodes, rewards_over_seeds, MAX_REWARD), f)
|
||||
|
||||
episode += 1
|
||||
episode = 0
|
||||
MAX_REWARD = 0
|
||||
i_seed += 1
|
||||
rewards_over_seeds.append(reward_over_episodes)
|
||||
agent = None # Reset the agent
|
||||
reward_over_episodes = None # Reset the reward_over_episodes
|
||||
# %%
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Reference in New Issue
Block a user