def main(): tf.set_random_seed(FLAGS.random_seed) np.random.seed(FLAGS.random_seed) random.seed(FLAGS.random_seed) graph = tf.Graph() # gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=FLAGS.gpu_memory_fraction) # tf_config = tf.ConfigProto(gpu_options=gpu_options) tf_config = tf.ConfigProto() tf_config.gpu_options.allow_growth = True sess = tf.Session(graph=graph, config=tf_config) network_config = { 'num_filters': [FLAGS.num_filters]*FLAGS.num_conv_layers, 'strides': [[1, 2, 2, 1]]*FLAGS.num_strides + [[1, 1, 1, 1]]*(FLAGS.num_conv_layers-FLAGS.num_strides), 'filter_size': FLAGS.filter_size, 'image_width': FLAGS.im_width, 'image_height': FLAGS.im_height, 'image_channels': FLAGS.num_channels, 'n_layers': FLAGS.num_fc_layers, 'layer_size': FLAGS.layer_size, 'initialization': FLAGS.init, } data_generator = DataGenerator() if FLAGS.train: state_idx = data_generator.state_idx else: state_idx = range(7) data_generator._dU = 7 img_idx = range(len(state_idx), len(state_idx)+FLAGS.im_height*FLAGS.im_width*FLAGS.num_channels) model = MIL(data_generator._dU, state_idx=state_idx, img_idx=img_idx, network_config=network_config) log_dir = os.path.join( FLAGS.data_path, 'logged_model') if FLAGS.train: with graph.as_default(): # train_image_tensors = data_generator.make_batch_tensor(network_config) # inputa = train_image_tensors[:, :FLAGS.number_of_shot*FLAGS.TimeFrame, :] # inputb = train_image_tensors[:, FLAGS.number_of_shot*FLAGS.TimeFrame:, :] inputa, inputb = data_generator.make_batch_tensor(network_config) train_input_tensors = {'inputa': inputa, 'inputb': inputb} # val_image_tensors = data_generator.make_batch_tensor(network_config, train=False) # inputa = val_image_tensors[:, :FLAGS.number_of_shot*FLAGS.TimeFrame, :] # inputb = val_image_tensors[:, FLAGS.number_of_shot*FLAGS.TimeFrame:, :] inputa, inputb = data_generator.make_batch_tensor(network_config, train=False) val_input_tensors = {'inputa': inputa, 'inputb': inputb} model.init_network(graph, input_tensors=train_input_tensors) model.init_network(graph, input_tensors=val_input_tensors, prefix='Validation_') else: model.init_network(graph, prefix='Testing') with graph.as_default(): # Set up saver. saver = tf.train.Saver(max_to_keep=10) # Initialize variables. init_op = tf.global_variables_initializer() sess.run(init_op, feed_dict=None) # Start queue runners (used for loading videos on the fly) tf.train.start_queue_runners(sess=sess) if FLAGS.resume: model_file = tf.train.latest_checkpoint(log_dir) print(model_file) #if FLAGS.restore_iter > 0: # model_file = model_file[:model_file.index('model')] + 'model_' + str(FLAGS.restore_iter) if model_file: with graph.as_default(): saver.restore(sess, model_file) if FLAGS.train: train(graph, model, saver, sess, data_generator, log_dir) else: robot_data_path = os.path.join( FLAGS.data_path, 'low_res_robot_data') load_one_shot_data_from_path(robot_data_path, data_generator, network_config) if FLAGS.experiment == 'reaching': control_robot(graph, model, data_generator, sess, 'reach', log_dir) else: control_robot(graph, model, data_generator, sess, 'push', log_dir)
def main(): tf.set_random_seed(FLAGS.random_seed) np.random.seed(FLAGS.random_seed) random.seed(FLAGS.random_seed) # Build up environment to prevent segfault if not FLAGS.train: if 'reach' in FLAGS.experiment: env = gym.make('ReacherMILTest-v1') ob = env.reset() # import pdb; pdb.set_trace() graph = tf.Graph() gpu_options = tf.GPUOptions( per_process_gpu_memory_fraction=FLAGS.gpu_memory_fraction) tf_config = tf.ConfigProto(gpu_options=gpu_options) tf_config.gpu_options.allow_growth = True sess = tf.Session(graph=graph, config=tf_config) network_config = { 'num_filters': [FLAGS.num_filters] * FLAGS.num_conv_layers, 'strides': [[1, 2, 2, 1]] * FLAGS.num_strides + [[1, 1, 1, 1]] * (FLAGS.num_conv_layers - FLAGS.num_strides), 'filter_size': FLAGS.filter_size, 'image_width': FLAGS.im_width, 'image_height': FLAGS.im_height, 'image_channels': FLAGS.num_channels, 'n_layers': FLAGS.num_fc_layers, 'layer_size': FLAGS.layer_size, 'initialization': FLAGS.init, } data_generator = DataGenerator() state_idx = data_generator.state_idx img_idx = range( len(state_idx), len(state_idx) + FLAGS.im_height * FLAGS.im_width * FLAGS.num_channels) # need to compute x_idx and img_idx from data_generator model = MIL(data_generator._dU, state_idx=state_idx, img_idx=img_idx, network_config=network_config) # TODO: figure out how to save summaries and checkpoints exp_string = FLAGS.experiment+ '.' + FLAGS.init + '_init.' + str(FLAGS.num_conv_layers) + '_conv' + '.' + str(FLAGS.num_strides) + '_strides' + '.' + str(FLAGS.num_filters) + '_filters' + \ '.' + str(FLAGS.num_fc_layers) + '_fc' + '.' + str(FLAGS.layer_size) + '_dim' + '.bt_dim_' + str(FLAGS.bt_dim) + '.mbs_'+str(FLAGS.meta_batch_size) + \ '.ubs_' + str(FLAGS.update_batch_size) + '.numstep_' + str(FLAGS.num_updates) + '.updatelr_' + str(FLAGS.train_update_lr) if FLAGS.clip: exp_string += '.clip_' + str(int(FLAGS.clip_max)) if FLAGS.conv_bt: exp_string += '.conv_bt' if FLAGS.all_fc_bt: exp_string += '.all_fc_bt' if FLAGS.fp: exp_string += '.fp' if FLAGS.learn_final_eept: exp_string += '.learn_ee_pos' if FLAGS.no_action: exp_string += '.no_action' if FLAGS.zero_state: exp_string += '.zero_state' if FLAGS.two_head: exp_string += '.two_heads' if FLAGS.two_arms: exp_string += '.two_arms' if FLAGS.temporal_conv_2_head: exp_string += '.1d_conv_act_' + str( FLAGS.temporal_num_layers) + '_' + str(FLAGS.temporal_num_filters) if FLAGS.temporal_conv_2_head_ee: exp_string += '_ee_' + str( FLAGS.temporal_num_layers_ee) + '_' + str( FLAGS.temporal_num_filters_ee) exp_string += '_' + str(FLAGS.temporal_filter_size) + 'x1_filters' if FLAGS.training_set_size != -1: exp_string += '.' + str(FLAGS.training_set_size) + '_trials' log_dir = FLAGS.log_dirs + '/' + exp_string # put here for now if FLAGS.train: data_generator.generate_batches(noisy=FLAGS.use_noisy_demos) with graph.as_default(): # train_image_tensors = data_generator.make_batch_tensor(network_config, restore_iter=FLAGS.restore_iter) train_image_tensors = data_generator.make_compare_batch_tensor( network_config, restore_iter=FLAGS.restore_iter) inputa = train_image_tensors[:, :FLAGS.update_batch_size * FLAGS.T, :] inputb = train_image_tensors[:, FLAGS.update_batch_size * FLAGS.T:(FLAGS.update_batch_size + 1) * FLAGS.T, :] inputc = train_image_tensors[:, (FLAGS.update_batch_size + 1) * FLAGS.T:, :] # train_input_tensors = {'inputa': inputa, 'inputb': inputb} train_input_tensors = { 'inputa': inputa, 'inputb': inputb, 'inputc': inputc } # val_image_tensors = data_generator.make_batch_tensor(network_config, restore_iter=FLAGS.restore_iter, train=False) # inputa = val_image_tensors[:, :FLAGS.update_batch_size*FLAGS.T, :] # inputb = val_image_tensors[:, FLAGS.update_batch_size*FLAGS.T:, :] # val_input_tensors = {'inputa': inputa, 'inputb': inputb} model.init_network(graph, input_tensors=train_input_tensors, restore_iter=FLAGS.restore_iter) # model.init_network(graph, input_tensors=val_input_tensors, restore_iter=FLAGS.restore_iter, prefix='Validation_') else: model.init_network(graph, prefix='Testing') with graph.as_default(): # Set up saver. saver = tf.train.Saver(max_to_keep=10) # Initialize variables. init_op = tf.global_variables_initializer() sess.run(init_op, feed_dict=None) # Start queue runners (used for loading videos on the fly) tf.train.start_queue_runners(sess=sess) if FLAGS.resume: model_file = tf.train.latest_checkpoint(log_dir) if FLAGS.restore_iter > 0: model_file = model_file[:model_file.index('model' )] + 'model_' + str( FLAGS.restore_iter) if model_file: ind1 = model_file.index('model') resume_itr = int(model_file[ind1 + 6:]) print("Restoring model weights from " + model_file) with graph.as_default(): saver.restore(sess, model_file) if FLAGS.train: train(graph, model, saver, sess, data_generator, log_dir, restore_itr=FLAGS.restore_iter) else: model_file = tf.train.latest_checkpoint(log_dir) if (FLAGS.begin_restore_iter != FLAGS.end_restore_iter): iter_index = FLAGS.begin_restore_iter while iter_index <= FLAGS.end_restore_iter: print('iter_index', iter_index) if FLAGS.restore_iter >= 0: model_file = model_file[:model_file.index( 'model')] + 'model_' + str(iter_index) if model_file: ind1 = model_file.index('model') resume_itr = int(model_file[ind1 + 6:]) print("Restoring model weights from " + model_file) # saver = tf.train.Saver() saver.restore(sess, model_file) if 'reach' in FLAGS.experiment: env = gym.make('ReacherMILTest-v1') env.reset() generate_test_demos(data_generator) evaluate_vision_reach(env, graph, model, data_generator, sess, exp_string, FLAGS.record_gifs, log_dir) # evaluate_rl_vision_reach(graph, data_generator, sess, exp_string, FLAGS.record_gifs, log_dirs) elif 'push' in FLAGS.experiment: evaluate_push(sess, graph, model, data_generator, exp_string, log_dir, FLAGS.demo_file + '/', save_video=FLAGS.record_gifs) iter_index += 100 else: if FLAGS.restore_iter > 0: model_file = model_file[:model_file. index('model')] + 'model_' + str( FLAGS.restore_iter) if model_file: ind1 = model_file.index('model') resume_itr = int(model_file[ind1 + 6:]) print("Restoring model weights from " + model_file) # saver = tf.train.Saver() saver.restore(sess, model_file) if 'reach' in FLAGS.experiment: env = gym.make('ReacherMILTest-v1') env.reset() generate_test_demos(data_generator) evaluate_vision_reach(env, graph, model, data_generator, sess, exp_string, FLAGS.record_gifs, log_dir) # evaluate_vision_reach(env, graph, data_generator, sess, exp_string, FLAGS.record_gifs, log_dir) # evaluate_rl_vision_reach(graph, data_generator, sess, exp_string, FLAGS.record_gifs, log_dirs) elif 'push' in FLAGS.experiment: evaluate_push(sess, graph, model, data_generator, exp_string, log_dir, FLAGS.demo_file + '/', save_video=FLAGS.record_gifs)
FLAGS.layer_size, 'initialization': FLAGS.init, 'temporal_conv_2_head_ee': FLAGS.temporal_conv_2_head_ee, } data_generator = DataGenerator() state_idx = data_generator.state_idx img_idx = range( len(state_idx), len(state_idx) + FLAGS.im_height * FLAGS.im_width * FLAGS.num_channels) model = MIL(data_generator._dU, state_idx=state_idx, img_idx=img_idx, network_config=network_config) model.init_network(graph, prefix='Testing') exp_string = FLAGS.experiment + '.' + FLAGS.init + '_init.' + str(FLAGS.num_conv_layers) + '_conv' + '.' + str(FLAGS.num_strides) + '_strides' + '.' + str(FLAGS.num_filters) + '_filters' + \ '.' + str(FLAGS.num_fc_layers) + '_fc' + '.' + str(FLAGS.layer_size) + '_dim' + '.bt_dim_' + str(FLAGS.bt_dim) + '.mbs_'+str(FLAGS.meta_batch_size) + \ '.ubs_' + str(FLAGS.update_batch_size) + '.numstep_' + str(FLAGS.num_updates) + '.updatelr_' + str(FLAGS.train_update_lr) if FLAGS.clip: exp_string += '.clip_' + str(int(FLAGS.clip_max)) if FLAGS.conv_bt: exp_string += '.conv_bt' if FLAGS.all_fc_bt: exp_string += '.all_fc_bt' if FLAGS.fp: exp_string += '.fp' if FLAGS.learn_final_eept: