import modules instead of importing packages (#76)

Author: yundiqian

compiler_opt/rl/agent_creators.py

@@ -19,7 +19,7 @@
 import gin
 import tensorflow as tf
 
-from tf_agents.agents import TFAgent
+from tf_agents.agents import tf_agent
 from tf_agents.agents.behavioral_cloning import behavioral_cloning_agent
 from tf_agents.agents.dqn import dqn_agent
 from tf_agents.agents.ppo import ppo_agent
@@ -29,10 +29,10 @@
 from compiler_opt.rl import constant_value_network
 
 
-def _create_behavioral_cloning_agent(time_step_spec: types.NestedTensorSpec,
-                                     action_spec: types.NestedTensorSpec,
-                                     preprocessing_layers: types.NestedLayer,
-                                     policy_network: types.Network) -> TFAgent:
+def _create_behavioral_cloning_agent(
+    time_step_spec: types.NestedTensorSpec, action_spec: types.NestedTensorSpec,
+    preprocessing_layers: types.NestedLayer,
+    policy_network: types.Network) -> tf_agent.TFAgent:
   """Creates a behavioral_cloning_agent."""
 
   network = policy_network(
@@ -48,7 +48,7 @@ def _create_behavioral_cloning_agent(time_step_spec: types.NestedTensorSpec,
 def _create_dqn_agent(time_step_spec: types.NestedTensorSpec,
                       action_spec: types.NestedTensorSpec,
                       preprocessing_layers: types.NestedLayer,
-                      policy_network: types.Network) -> TFAgent:
+                      policy_network: types.Network) -> tf_agent.TFAgent:
   """Creates a dqn_agent."""
   network = policy_network(
       time_step_spec.observation,
@@ -62,7 +62,7 @@ def _create_dqn_agent(time_step_spec: types.NestedTensorSpec,
 def _create_ppo_agent(time_step_spec: types.NestedTensorSpec,
                       action_spec: types.NestedTensorSpec,
                       preprocessing_layers: types.NestedLayer,
-                      policy_network: types.Network) -> TFAgent:
+                      policy_network: types.Network) -> tf_agent.TFAgent:
   """Creates a ppo_agent."""
 
   actor_network = policy_network(
@@ -87,7 +87,7 @@ def create_agent(agent_name: constant.AgentName,
                  action_spec: types.NestedTensorSpec,
                  preprocessing_layer_creator: Callable[[types.TensorSpec],
                                                        tf.keras.layers.Layer],
-                 policy_network: types.Network) -> TFAgent:
+                 policy_network: types.Network) -> tf_agent.TFAgent:
   """Creates a tfa.agents.TFAgent object.
 
   Args:

compiler_opt/rl/policy_saver.py

@@ -18,7 +18,7 @@
 import os
 
 import tensorflow as tf
-from tf_agents.policies import TFPolicy
+from tf_agents.policies import tf_policy
 from tf_agents.policies import policy_saver
 
 from typing import Dict, Tuple
@@ -77,14 +77,14 @@ class PolicySaver(object):
   ```
   """
 
-  def __init__(self, policy_dict: Dict[str, TFPolicy]):
+  def __init__(self, policy_dict: Dict[str, tf_policy.TFPolicy]):
     """Initialize the PolicySaver object.
 
     Args:
       policy_dict: A dict mapping from policy name to policy.
     """
     self._policy_saver_dict: Dict[str, Tuple[
-        policy_saver.PolicySaver, TFPolicy]] = {
+        policy_saver.PolicySaver, tf_policy.TFPolicy]] = {
             policy_name: (policy_saver.PolicySaver(
                 policy, batch_size=1, use_nest_path_signatures=False), policy)
             for policy_name, policy in policy_dict.items()

compiler_opt/rl/train_bc.py

@@ -30,8 +30,8 @@
 from compiler_opt.rl import registry
 from compiler_opt.rl import trainer
 
-from tf_agents.agents import TFAgent
-from tf_agents.policies import TFPolicy
+from tf_agents.agents import tf_agent
+from tf_agents.policies import tf_policy
 
 from typing import Dict
 
@@ -63,13 +63,12 @@ def train_eval(agent_name=constant.AgentName.BEHAVIORAL_CLONE,
   preprocessing_layer_creator = problem_config.get_preprocessing_layer_creator()
 
   # Initialize trainer and policy saver.
-  tf_agent: TFAgent = agent_creators.create_agent(agent_name, time_step_spec,
-                                                  action_spec,
-                                                  preprocessing_layer_creator)
-  llvm_trainer = trainer.Trainer(root_dir=root_dir, agent=tf_agent)
-  policy_dict: Dict[str, TFPolicy] = {
-      'saved_policy': tf_agent.policy,
-      'saved_collect_policy': tf_agent.collect_policy,
+  agent: tf_agent.TFAgent = agent_creators.create_agent(
+      agent_name, time_step_spec, action_spec, preprocessing_layer_creator)
+  llvm_trainer = trainer.Trainer(root_dir=root_dir, agent=agent)
+  policy_dict: Dict[str, tf_policy.TFPolicy] = {
+      'saved_policy': agent.policy,
+      'saved_collect_policy': agent.collect_policy,
   }
   saver = policy_saver.PolicySaver(policy_dict=policy_dict)
 

compiler_opt/rl/train_locally.py

@@ -25,7 +25,7 @@
 from absl import logging
 import gin
 import tensorflow as tf
-from tf_agents.agents import TFAgent
+from tf_agents.agents import tf_agent
 from tf_agents.system import system_multiprocessing as multiprocessing
 from typing import List
 
@@ -78,9 +78,8 @@ def train_eval(agent_name=constant.AgentName.PPO,
   preprocessing_layer_creator = problem_config.get_preprocessing_layer_creator()
 
   # Initialize trainer and policy saver.
-  tf_agent: TFAgent = agent_creators.create_agent(agent_name, time_step_spec,
-                                                  action_spec,
-                                                  preprocessing_layer_creator)
+  agent: tf_agent.TFAgent = agent_creators.create_agent(
+      agent_name, time_step_spec, action_spec, preprocessing_layer_creator)
   # create the random network distillation object
   random_network_distillation = None
   if use_random_network_distillation:
@@ -91,13 +90,13 @@ def train_eval(agent_name=constant.AgentName.PPO,
 
   llvm_trainer = trainer.Trainer(
       root_dir=root_dir,
-      agent=tf_agent,
+      agent=agent,
       random_network_distillation=random_network_distillation,
       warmstart_policy_dir=warmstart_policy_dir)
 
   policy_dict = {
-      'saved_policy': tf_agent.policy,
-      'saved_collect_policy': tf_agent.collect_policy,
+      'saved_policy': agent.policy,
+      'saved_collect_policy': agent.collect_policy,
   }
   saver = policy_saver.PolicySaver(policy_dict=policy_dict)
 

compiler_opt/rl/trainer.py

@@ -21,7 +21,7 @@
 import gin
 import tensorflow as tf
 from compiler_opt.rl import random_net_distillation
-from tf_agents.agents import TFAgent
+from tf_agents.agents import tf_agent
 from tf_agents.policies import policy_loader
 
 from tf_agents.utils import common as common_utils
@@ -47,7 +47,7 @@ class Trainer(object):
   def __init__(
       self,
       root_dir: str,
-      agent: TFAgent,
+      agent: tf_agent.TFAgent,
       random_network_distillation: Optional[
           random_net_distillation.RandomNetworkDistillation] = None,
       warmstart_policy_dir: Optional[str] = None,