build framework for training

agent.py

@@ -74,8 +74,8 @@ class Policy(nn.Module):
   def __init__(self, 
                actor,
                critic,
-               obs_space, 
-               action_space,
+               obs_space=None, 
+               action_space=None,
                optimizer=None,
                lr=2.5e-4,
                weight_decay=5e-4,

env/env.py

@@ -25,9 +25,8 @@ class RSMTEnv(gym.Env):
     })
 
     self.action_space = spaces.Dict({
-      "u": spaces.Discrete(num_nodes),
-      "w": spaces.Discrete(num_nodes),
-      "s": spaces.Discrete(2),
+      "v": spaces.Discrete(num_nodes),
+      "h": spaces.Discrete(num_nodes),
     })
 
     self.mask_unvisited = None
@@ -50,3 +49,7 @@ class RSMTEnv(gym.Env):
   
   def close(self):
     pass
+
+def create_RSMTEnv(*args, **kwarg):
+  env = RSMTEnv(*args, **kwarg)
+  return env
\ No newline at end of file

test.py

@@ -2,7 +2,7 @@ import torch
 import pdb
 
 from agent import Actor, Critic, Policy
-from env.env import RSMTEnv
+from env.env import create_RSMTEnv
 
 def update_mask(mask, batch_size, node_visited):
   for i in range(batch_size):
@@ -13,7 +13,7 @@ batch_size = 4
 num_nodes = 8
 nodes = torch.randn(batch_size, num_nodes, 2)
 
-env = RSMTEnv(num_nodes=num_nodes, pos_l=0, pos_h=100)
+env = create_RSMTEnv(num_nodes=num_nodes, pos_l=0, pos_h=100)
 
 actor_net = Actor()
 critic_net = Critic()

train.py

+import torch
+
+from agent import Actor, Critic, Policy
+from env.env import create_RSMTEnv
+import gym
+
+import pdb
+
+batch_size = 4
+num_nodes = 8
+max_episode_num = 40000 # 40k
+
+pos_l = 0.0
+pos_h = 100.0
+
+env_fns = [lambda: create_RSMTEnv(num_nodes=num_nodes, 
+                                  pos_l=pos_l, 
+                                  pos_h=pos_h)
+           for i in range (batch_size)
+          ]
+envs = gym.vector.AsyncVectorEnv(env_fns)
+
+actor_net = Actor()
+critic_net = Critic()
+model = Policy(actor=actor_net, 
+               critic=critic_net, 
+              )
+
+nodes_xy, mask_visited = envs.reset()
+for episode in range(max_episode_num):
+  v, h = model.act(nodes_xy, mask_visited=mask_visited)
+  # v/h is a tensor of shape [#batch_size] on model.device
+  mask_visited, rewards, done, _ = envs.step(v, h)
+  if done:
+    model.learn(nodes_xy, rewards)
+    nodes_xy, mask_visited = envs.reset()
+