Refactoring agent structure to be more method oriented.

digideep/agent/__init__.py

@@ -1,3 +0,0 @@
-from .ppo import PPO
-from .ddpg import DDPG
-from .sac import SAC

digideep/agent/policy/soft_stochastic/__init__.py → digideep/agent/ddpg/__init__.py

@@ -1 +1,2 @@
+from .agent import Agent
 from .policy import Policy

digideep/agent/ddpg.py → digideep/agent/ddpg/agent.py

@@ -7,21 +7,19 @@
 import torch.nn.functional as F
 import torch.utils.data
 
-from digideep.agent.samplers.ddpg import sampler_re
-# from digideep.agent.samplers.common import check_shape
-
 from digideep.utility.toolbox import get_class
 from digideep.utility.logging import logger
 from digideep.utility.profiling import KeepTime
 from digideep.utility.monitoring import monitor
 
-
-from .base import AgentBase
-from digideep.agent.policy.deterministic import Policy
+# from digideep.agent.sampler_common import check_shape
+from digideep.agent.agent_base import AgentBase
+from .sampler import sampler_re
+from .policy import Policy
 
 # torch.utils.backcompat.broadcast_warning.enabled = True
 
-class DDPG(AgentBase):
+class Agent(AgentBase):
     """This is an implementation of the Deep Deterministic Policy Gradient (`DDPG <https://arxiv.org/abs/1509.02971>`_) method.
     
     Args:
@@ -46,7 +44,7 @@ class DDPG(AgentBase):
     """
 
     def __init__(self, session, memory, **params):
-        super(DDPG, self).__init__(session, memory, **params)
+        super(Agent, self).__init__(session, memory, **params)
 
         self.device = self.session.get_device()
 

digideep/agent/policy/deterministic/policy.py → digideep/agent/ddpg/policy.py

@@ -6,8 +6,8 @@
 from digideep.utility.toolbox import get_class
 from digideep.utility.logging import logger
 
-from digideep.agent.policy.base import PolicyBase
-from digideep.agent.policy.common import Averager
+from digideep.agent.policy_base import PolicyBase
+from digideep.agent.policy_common import Averager
 
 from copy import deepcopy
 

digideep/agent/samplers/ddpg.py → digideep/agent/ddpg/sampler.py

@@ -1,17 +1,16 @@
 import numpy as np
 import warnings
 
-from .common import Compose, flatten_memory_to_train_key, get_memory_params, check_nan, check_shape, check_stats, print_line
-from .common import flatten_first_two
-from digideep.utility.logging import logger
+from digideep.agent.sampler_common import Compose, flatten_memory_to_train_key, get_memory_params, check_nan, check_shape, check_stats, print_line
+from digideep.agent.sampler_common import flatten_first_two
 
+from digideep.utility.logging import logger
 from digideep.utility.profiling import KeepTime
 
 def get_sample_memory(buffer, info):
     """Sampler function for DDPG-like algorithms where we want to sample data from an experience replay buffer.
 
-    This function does not sample from final steps where mask equals ``0`` (as they don't have subsequent observations.)
-    This function adds the following key to the memory:
+    This function adds the following key to the buffer:
     
     * ``/observations_2``
 

digideep/agent/policy/deterministic/__init__.py → digideep/agent/ppo/__init__.py

@@ -1 +1,2 @@
+from .agent import Agent
 from .policy import Policy

digideep/agent/ppo.py → digideep/agent/ppo/agent.py

@@ -7,15 +7,15 @@
 import torch.nn.functional as F
 import torch.utils.data
 
-from digideep.agent.samplers.ppo import sampler_ff, sampler_rn
+from .sampler import sampler_ff, sampler_rn
 # from digideep.agent.samplers.common import check_shape
 
 from digideep.utility.toolbox import get_class
 from digideep.utility.logging import logger
 from digideep.utility.profiling import KeepTime
 from digideep.utility.monitoring import monitor
 
-from .base import AgentBase
+from digideep.agent.agent_base import AgentBase
 
 
 # def update_linear_schedule(optimizer, epoch, total_num_epochs, initial_lr):
@@ -25,7 +25,7 @@
 #         param_group['lr'] = lr
 
 
-class PPO(AgentBase):
+class Agent(AgentBase):
     """The implementation of the Proximal Policy Optimization (`PPO <https://arxiv.org/abs/1707.06347>`_) method.
 
     Args:
@@ -50,7 +50,7 @@ class PPO(AgentBase):
         
     """
     def __init__(self, session, memory, **params):
-        super(PPO, self).__init__(session, memory, **params)
+        super(Agent, self).__init__(session, memory, **params)
 
         self.device = self.session.get_device()