add setting

train_dreamer_add_class_gen.py

@@ -6,15 +6,16 @@
 import torch.nn.parallel
 import torch.backends.cudnn as cudnn
 import torch.optim as optim
+from utils_dreamer import parse_arguments
 opt = parse_arguments()
-if opt.task == 'max':
+if opt.setting == 'max':
     from architectures_dreamer_max import Sender, ReceiverOnestep, ValueModel, Players
+elif opt.setting == 'retrieve':
+    from architectures_retrieve_5stroke import Sender, ReceiverOnestep, ValueModel, Players
 else:
     from architectures_dreamer_v2 import Sender, ReceiverOnestep, ValueModel, Players
-from torch.autograd import Variable
 import pdb
 from reinforce_one_stroke import *
-from utils_dreamer import parse_arguments
 from utils_dreamer import get_batch_random as get_batch_train_func
 import numpy as np
 import pickle

utils_dreamer.py

@@ -28,9 +28,11 @@ def parse_arguments():
     parser.add_argument(
         '--data_root', default='./data/output/', help='data root folder')
     parser.add_argument(
-        '--sender_path', default='../ddpg_comm/pretrained/actor_one_stroke.pkl', help='pretrained folder')
+        '--sender_path', default='./pretrained/actor_one_stroke.pkl', help='pretrained folder')
     parser.add_argument(
         '--resume_path', default=None, help='pretrained folder')
+    parser.add_argument(
+        '--setting', default='complete', help='game settings')
     parser.add_argument('--max_step', type=int,
                         help='number of drawing steps', default=10)
     parser.add_argument('--num_stroke', type=int,