Fix spurious AgentConfig member reference (#277)

mtrofin · web-flow · commit b82deb3e70c9 · 2023-07-12T15:55:59.000-07:00
diff --git a/compiler_opt/rl/distributed/ppo_collect_lib.py b/compiler_opt/rl/distributed/ppo_collect_lib.py
@@ -124,7 +124,7 @@ def collect(corpus_path: str, replay_buffer_server_address: str,
   agent_cfg = agent_config.DistributedPPOAgentConfig(
       time_step_spec=time_step_spec, action_spec=action_spec)
   agent = agent_config.create_agent(
-      agent_cfg.agent,
+      agent_cfg,
       preprocessing_layer_creator=problem_config
       .get_preprocessing_layer_creator())
 
diff --git a/compiler_opt/rl/distributed/ppo_eval_lib.py b/compiler_opt/rl/distributed/ppo_eval_lib.py
@@ -60,7 +60,7 @@ def evaluate(root_dir: str, corpus_path: str,
   agent_cfg = agent_config.DistributedPPOAgentConfig(
       time_step_spec=time_step_spec, action_spec=action_spec)
   agent = agent_config.create_agent(
-      agent_cfg.agent,
+      agent_cfg,
       preprocessing_layer_creator=problem_config
       .get_preprocessing_layer_creator())
 
diff --git a/compiler_opt/rl/train_locally.py b/compiler_opt/rl/train_locally.py
@@ -82,7 +82,7 @@ def train_eval(worker_manager_class=LocalWorkerPoolManager,
   agent_cfg = agent_config_type(
       time_step_spec=time_step_spec, action_spec=action_spec)
   agent: tf_agent.TFAgent = agent_config.create_agent(
-      agent_cfg.agent, preprocessing_layer_creator=preprocessing_layer_creator)
+      agent_cfg, preprocessing_layer_creator=preprocessing_layer_creator)
   # create the random network distillation object
   random_network_distillation = None
   if use_random_network_distillation:
