fix minor (#62)

Author: ramanuzan

jorldy/core/agent/ddpg.py

@@ -12,6 +12,7 @@
 
 
 class DDPG(BaseAgent):
+    action_type = "continuous"
     """Deep deterministic policy gradient (DDPG) agent.
 
     Args:
@@ -65,7 +66,6 @@ def __init__(
             if device
             else torch.device("cuda" if torch.cuda.is_available() else "cpu")
         )
-        self.action_type = "continuous"
 
         self.actor = Network(
             actor, state_size, action_size, D_hidden=hidden_size, head=head

jorldy/core/agent/dqn.py

@@ -12,6 +12,7 @@
 
 
 class DQN(BaseAgent):
+    action_type = "discrete"
     """DQN agent.
 
     Args:
@@ -57,7 +58,6 @@ def __init__(
         num_workers=1,
         **kwargs,
     ):
-
         self.device = (
             torch.device(device)
             if device

jorldy/core/agent/ppo.py

@@ -25,6 +25,7 @@ class PPO(REINFORCE):
 
     def __init__(
         self,
+        network="discrete_policy_value",
         batch_size=32,
         n_step=128,
         n_epoch=3,
@@ -36,7 +37,7 @@ def __init__(
         num_workers=1,
         **kwargs,
     ):
-        super(PPO, self).__init__(**kwargs)
+        super(PPO, self).__init__(network=network, **kwargs)
 
         self.batch_size = batch_size
         self.n_step = n_step

jorldy/core/agent/vmpo.py

@@ -31,6 +31,7 @@ class VMPO(REINFORCE):
 
     def __init__(
         self,
+        network="discrete_policy_value",
         optim_config={"name": "adam"},
         batch_size=32,
         n_step=128,
@@ -49,7 +50,11 @@ def __init__(
         alpha_sigma=1.0,
         **kwargs,
     ):
-        super(VMPO, self).__init__(optim_config=optim_config, **kwargs)
+        super(VMPO, self).__init__(
+            network=network,
+            optim_config=optim_config,
+            **kwargs,
+        )
 
         self.batch_size = batch_size
         self.n_step = n_step