Moves fsm into rebar

andyljones · Jul 7, 2020 · 1a1d078 · 1a1d078
1 parent 72dc820
commit 1a1d078
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/onedee/fsm.py → rebar/fsm.py b/onedee/fsm.py → rebar/fsm.py
@@ -1,11 +1,22 @@
 import numpy as np
 import torch
 from rebar import arrdict, dotdict
-from . import spaces
 import pandas as pd
 
 __all__ = ['dataframe']
 
+class MultiVector:
+
+    def __init__(self, n_agents, dim):
+        super().__init__()
+        self.shape = (n_agents, dim)
+
+class MultiDiscrete:
+
+    def __init__(self, n_agents, n_actions):
+        super().__init__()
+        self.shape = (n_agents, n_actions)
+
 def _dataframe(traj):
     if isinstance(traj, dict):
         return [([k] + kk, vv) for k, v in traj.items() for kk, vv in _dataframe(v)]
@@ -32,8 +43,8 @@ def __init__(self, n_envs, fsm, device='cuda'):
 
         self._token = torch.full((self.n_envs,), -1, dtype=torch.long, device=device)
 
-        self.observation_space = spaces.MultiVector(1, fsm.d_obs) if fsm.d_obs else spaces.MultiEmpty()
-        self.action_space = spaces.MultiDiscrete(1, fsm.n_actions)
+        self.observation_space = MultiVector(1, fsm.d_obs) if fsm.d_obs else spaces.MultiEmpty()
+        self.action_space = MultiDiscrete(1, fsm.n_actions)
 
     def _reset(self, reset):
         if reset.any():