Update ACVP

main.py

@@ -302,8 +302,8 @@ def action_conditional_video_prediction():
     # a2c_pixel_atari(game)
 
     # Generate a dataset with the trained model
-    a2c_model_file = './data/A2CAgent-vanilla-model-%s.bin' % (game)
-    generate_dataset(game, a2c_model_file, prefix)
+    # a2c_model_file = './data/A2CAgent-vanilla-model-%s.bin' % (game)
+    # generate_dataset(game, a2c_model_file, prefix)
 
     # Train the action conditional video prediction model
     acvp_train(game, prefix)
@@ -334,7 +334,7 @@ if __name__ == '__main__':
     # ddpg_continuous()
     # ppo_continuous()
 
-    action_conditional_video_prediction()
+    # action_conditional_video_prediction()
 
     # plot()
 

model/action_conditional_video_prediction.py

@@ -4,6 +4,8 @@
 # declaration at the top                                              #
 #######################################################################
 
+__all__ = ['acvp_train']
+
 import torch
 from torch.autograd import Variable
 import torch.nn as nn

model/dataset.py

@@ -50,6 +50,7 @@ def generate_dataset(game, a2c_model, prefix):
     config.gradient_clip = 0.5
     config.logger = Logger('./log', logger, skip=True)
     agent = A2CAgent(config)
+    agent.close()
 
     agent.load(a2c_model)
     task = PixelAtari(game, frame_skip=4, history_length=4, log_dir=None, dataset=True)