Merge pull request #93 from UT-Austin-RPL/discret_dist_type

`discrete_dist_type`

Author: jakegrigsby

amago/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "3.1.1"
+__version__ = "3.1.2"
 
 from .experiment import Experiment
 from .agent import Agent

amago/nets/actor_critic.py

@@ -34,6 +34,7 @@ def __init__(
         discrete: bool,
         gammas: torch.Tensor,
         continuous_dist_type: Type[PolicyOutput],
+        discrete_dist_type: Type[PolicyOutput] = Discrete,
     ):
         super().__init__()
         self.state_dim = state_dim
@@ -42,7 +43,7 @@ def __init__(
         self.gammas = gammas
         self.num_gammas = len(self.gammas)
         # determine policy output
-        dist_type = Discrete if self.discrete else continuous_dist_type
+        dist_type = discrete_dist_type if self.discrete else continuous_dist_type
         self.policy_dist = dist_type(d_action=self.action_dim)
         assert isinstance(self.policy_dist, PolicyOutput)
         assert self.policy_dist.is_discrete == self.discrete
@@ -116,13 +117,15 @@ def __init__(
         activation: str = "leaky_relu",
         dropout_p: float = 0.0,
         continuous_dist_type: Type[PolicyOutput] = TanhGaussian,
+        discrete_dist_type: Type[PolicyOutput] = Discrete,
     ):
         super().__init__(
             state_dim=state_dim,
             action_dim=action_dim,
             discrete=discrete,
             gammas=gammas,
             continuous_dist_type=continuous_dist_type,
+            discrete_dist_type=discrete_dist_type,
         )
         # build base network
         self.base = MLP(
@@ -184,13 +187,15 @@ def __init__(
         normalization: str = "layer",
         dropout_p: float = 0.0,
         continuous_dist_type: Type[PolicyOutput] = TanhGaussian,
+        discrete_dist_type: Type[PolicyOutput] = Discrete,
     ):
         super().__init__(
             state_dim=state_dim,
             action_dim=action_dim,
             discrete=discrete,
             gammas=gammas,
             continuous_dist_type=continuous_dist_type,
+            discrete_dist_type=discrete_dist_type,
         )
         self.inp = MLP(
             d_inp=state_dim,

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="amago",
-    version="3.1.1",
+    version="3.1.2",
     author="Jake Grigsby",
     author_email="grigsby@cs.utexas.edu",
     license="MIT",