Fix agent testing

verlab · Feb 17, 2021 · afff8c6 · afff8c6
1 parent 0aaae03
commit afff8c6
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/rl_fast_forward/REINFORCE/agent.py b/rl_fast_forward/REINFORCE/agent.py
@@ -161,6 +161,7 @@ def train(self, envs, dataset_name, n_epochs, model_path=None):
  avg_epoch_losses.append(epoch_losses.detach().cpu())
  avg_critic_losses.append(critic_losses.detach().cpu())
 
+ self.writer.add_scalar('Rewards_{}/_overall_avg_reward'.format(dataset_name), np.mean([scores[exp_key] for exp_key in exp_keys]), i_epoch)
 
  self.optimizer.zero_grad()
  epoch_losses.backward() # Computes the derivative of loss with respect to theta (dLoss/dTheta)

diff --git a/rl_fast_forward/main.py b/rl_fast_forward/main.py
@@ -133,6 +133,7 @@ def test(args):
  test_set = [exp for exp in experiments if Experiment2VideoMapping(exp).split_type in ['validation', 'test']]
 
  test_envs = {exp_key: VideoEnvironment(args.semantic_encoder_model_filename, args.user_document_filename, experiment_name=exp_key, batch_size=args.batch_size) for exp_key in test_set}
+ env = test_envs[test_set[0]]
 
  json_sf = {'info': {'version': 'v1.1_{}'.format(dt.now().strftime('%Y%m%d_%H%M%S')), 'dataset': args.dataset}, 'data': {}}