internal merge of PR #1277

PiperOrigin-RevId: 224250704

Author: koz4k

tensor2tensor/models/research/rl.py

@@ -61,14 +61,6 @@ def ppo_base_v1():
   return hparams
 
 
-#@registry.register_hparams
-#def ppo_continuous_action_base():
-#  hparams = ppo_base_v1()
-#  hparams.add_hparam("policy_network", feed_forward_gaussian_fun)
-#  hparams.add_hparam("policy_network_params", "basic_policy_parameters")
-#  return hparams
-
-
 @registry.register_hparams
 def basic_policy_parameters():
   wrappers = None
@@ -158,7 +150,7 @@ def get_policy(observations, hparams, action_space):
   """Get a policy network.
 
   Args:
-    observations
+    observations: observations
     hparams: parameters
     action_space: action space
 

tensor2tensor/rl/dopamine_connector.py

@@ -19,10 +19,9 @@
 from __future__ import division
 from __future__ import print_function
 
-from copy import copy
+import copy
+
 from dopamine.agents.dqn import dqn_agent
-from dopamine.agents.dqn.dqn_agent import NATURE_DQN_OBSERVATION_SHAPE
-from dopamine.agents.dqn.dqn_agent import NATURE_DQN_STACK_SIZE
 from dopamine.atari import run_experiment
 from dopamine.replay_memory import circular_replay_buffer
 from dopamine.replay_memory.circular_replay_buffer import OutOfGraphReplayBuffer
@@ -45,7 +44,7 @@
 
 
 class ResizeObservation(gym.ObservationWrapper):
-  """ TODO(konradczechowski): Add doc-string."""
+  """TODO(konradczechowski): Add doc-string."""
 
   def __init__(self, env, size=84):
     """Based on WarpFrame from openai baselines atari_wrappers.py.
@@ -91,7 +90,7 @@ def step(self, action):
 
 
 class _DQNAgent(dqn_agent.DQNAgent):
-  """ Modify dopamine DQNAgent to match our needs.
+  """Modify dopamine DQNAgent to match our needs.
 
   Allow passing batch_size and replay_capacity to ReplayBuffer, allow not using
   (some of) terminal episode transitions in training.
@@ -107,8 +106,8 @@ def __init__(self, replay_capacity, batch_size, generates_trainable_dones,
   def _build_replay_buffer(self, use_staging):
     """Build WrappedReplayBuffer with custom OutOfGraphReplayBuffer."""
     replay_buffer_kwargs = dict(
-        observation_shape=NATURE_DQN_OBSERVATION_SHAPE,
-        stack_size=NATURE_DQN_STACK_SIZE,
+        observation_shape=dqn_agent.NATURE_DQN_OBSERVATION_SHAPE,
+        stack_size=dqn_agent.NATURE_DQN_STACK_SIZE,
         replay_capacity=self._replay_capacity,
         batch_size=self._batch_size,
         update_horizon=self.update_horizon,
@@ -127,7 +126,7 @@ def _build_replay_buffer(self, use_staging):
 
 
 class _OutOfGraphReplayBuffer(OutOfGraphReplayBuffer):
-  """ Replay not sampling artificial_terminal transition.
+  """Replay not sampling artificial_terminal transition.
 
   Adds to stored tuples 'artificial_done' field (as last ReplayElement).
   When sampling, ignores tuples for which artificial_done is True.
@@ -238,7 +237,7 @@ def _get_optimizer(params):
 
 
 class DQNLearner(PolicyLearner):
-  """ Interface for learning dqn implemented in dopamine."""
+  """Interface for learning dqn implemented in dopamine."""
 
   def __init__(self, frame_stack_size, base_event_dir, agent_model_dir):
     super(DQNLearner, self).__init__(frame_stack_size, base_event_dir,
@@ -296,7 +295,7 @@ def train(self,
     if num_env_steps is None:
       num_env_steps = hparams.num_frames
 
-    hparams = copy(hparams)
+    hparams = copy.copy(hparams)
     hparams.set_hparam(
         "agent_epsilon_eval", min(hparams.agent_epsilon_eval * sampling_temp, 1)
     )
@@ -318,7 +317,7 @@ def evaluate(self, env_fn, hparams, sampling_temp):
     target_iterations = 0
     training_steps_per_iteration = 0
 
-    hparams = copy(hparams)
+    hparams = copy.copy(hparams)
     hparams.set_hparam(
         "agent_epsilon_eval", min(hparams.agent_epsilon_eval * sampling_temp, 1)
     )

tensor2tensor/rl/policy_learner.py

@@ -41,6 +41,7 @@ def train(
       eval_env_fn=None,
       report_fn=None
   ):
+    """Train."""
     # TODO(konradczechowski): pass name_scope instead of epoch?
     # TODO(konradczechowski): move 'simulated' to  batch_env
     raise NotImplementedError()