Hotfix policy shape

sharif1093 · Mar 8, 2019 · 7a97d31 · 7a97d31
1 parent 1b7ace3
commit 7a97d31
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/digideep/environment/make_environment.py b/digideep/environment/make_environment.py
@@ -45,7 +45,7 @@ def space2config(S):
     # S.__class__.__name__: "Discrete" / "Box"
     if isinstance(S, spaces.Discrete):
         typ = S.__class__.__name__
-        dim = S.n
+        dim = np.int32(S.n)
         lim = (np.nan, np.nan) # Discrete Spaces do not have high/low
         config = {"typ":typ, "dim":dim, "lim":lim}
     elif isinstance(S, spaces.Box):

diff --git a/digideep/policy/stochastic/policy.py b/digideep/policy/stochastic/policy.py
@@ -51,8 +51,8 @@ def __init__(self, device, obs_space, act_space, modelname, modelargs):
 
         # TODO: For discrete actions, `act_space["dim"][0]` works. It works for constinuous actions as well.
         #       Even for discrete actions `np.isscalar(act_space["dim"])` returns False.
-        # num_outputs = act_space["dim"] if np.isscalar(act_space["dim"]) else act_space["dim"][0]
-        num_outputs = act_space["dim"].item() if len(act_space["dim"].shape)==0 else act_space["dim"][0]
+        num_outputs = act_space["dim"] if np.isscalar(act_space["dim"]) else act_space["dim"][0]
+        # num_outputs = act_space["dim"].item() if len(act_space["dim"].shape)==0 else act_space["dim"][0]
         if act_space["typ"] == "Discrete":
             print("Discrete is recognized and num_outputs=", num_outputs)
             self.model["dist"] = Categorical(num_inputs=modelargs["output_size"], num_outputs=num_outputs)