Ejemplo n.º 1
0
def main():
    parser = argparse.ArgumentParser(
        description='Evolutionary training of controller')
    parser.add_argument('--env_name',
                        nargs='?',
                        default='CarRacing-v0',
                        help='Environment to use (default=CarRacing-v0)')
    parser.add_argument(
        '--n_rollouts',
        type=int,
        default=1,
        help='How many rollouts to perform when evaluating (default=1)')
    parser.add_argument('--n_generations',
                        type=int,
                        default=300,
                        help='Number of generations to train (default=300)')
    parser.add_argument('--latent_dim',
                        type=int,
                        default=32,
                        help='Dimension of latent space (default=32)')
    parser.add_argument('--seq_len',
                        type=int,
                        default=10,
                        help='Length of sequences for learning (default=10)')
    parser.add_argument('--action_dim',
                        type=int,
                        default=3,
                        help='Dimension of action space (default=3)')
    parser.add_argument('--rnn_hidden_dim',
                        nargs='?',
                        type=int,
                        default=256,
                        help='Dimension of RNN hidden state (default=256)')
    parser.add_argument(
        '--n_gaussians',
        type=int,
        default=5,
        help='Number of gaussians for the Mixture Density Network (default=5)')
    parser.add_argument(
        '--pop_size',
        type=int,
        default=64,
        help='Population size for evolutionary search (default=64)')
    parser.add_argument(
        '--n_workers',
        type=int,
        default=32,
        help='Number of workers for parallel processing (default=32)')
    parser.add_argument('--vae_fname', help='VAE model file name')
    parser.add_argument('--rnn_fname', nargs='?', help='RNN model file name')
    parser.add_argument(
        '--eval_interval',
        nargs='?',
        default=15,
        type=int,
        help='After how many generation to evaluate best params (default=15)')
    args = parser.parse_args()

    device = torch.device('cpu')

    # Load the VAE model from file.
    vae = VAE(latent_dim=args.latent_dim)
    vae.load_state_dict(
        torch.load(os.path.join(DATA_DIR, 'vae', args.vae_fname),
                   map_location={'cuda:0':
                                 'cpu'}))  # Previously trained on GPU.
    vae.to(device)

    # TODO: add identity/None RNN for dealing with the below?
    if args.rnn_fname is not None:  # Use memory module.
        # Load the MDNRNN model from file.
        mdnrnn = MDNRNN(action_dim=args.action_dim,
                        hidden_dim=args.rnn_hidden_dim,
                        latent_dim=args.latent_dim,
                        n_gaussians=args.n_gaussians)
        mdnrnn.load_state_dict(
            torch.load(os.path.join(DATA_DIR, 'rnn', args.rnn_fname),
                       map_location={'cuda:0':
                                     'cpu'}))  # Previously trained on GPU.
        mdnrnn.to(device)
    else:  # TODO: hacky, but dunno how to have default value for dim and pass it later without too many ifs. Fix?
        args.rnn_hidden_dim = 0
        mdnrnn = None

    # Set up controller model.
    agent = ControllerAgent(vae=vae,
                            v_dim=args.latent_dim,
                            action_dim=args.action_dim,
                            rnn=mdnrnn,
                            m_dim=args.rnn_hidden_dim)

    # Set up evolutionary strategy optimizer.
    with suppress_stdout(
    ):  # Suppress evolutionary strategy optimizer creation message.
        es = CMAES(
            num_params=param_count(agent.controller),
            sigma_init=0.1,  # initial standard deviation
            popsize=args.pop_size)

    # Set up multiprocessing.
    pool = mp.Pool(processes=args.n_workers)

    # Create results folder.
    dir_name = datetime.datetime.today().isoformat() + '_' + str(
        args.rnn_hidden_dim)
    os.makedirs(os.path.join(RESULTS_DIR, 'controller', dir_name))

    # TODO: add antithetic?
    for i in range(1, args.n_generations + 1):
        start_time = datetime.datetime.now()

        # Create a set of candidate specimens.
        specimens = es.ask()

        # Evaluate the fitness of candidate specimens.
        func = partial(evaluate,
                       env_name=args.env_name,
                       vae=vae,
                       rnn=mdnrnn,
                       v_dim=args.latent_dim,
                       action_dim=args.action_dim,
                       m_dim=args.rnn_hidden_dim,
                       n_rollouts=args.n_rollouts)
        fitness_list = np.array(pool.map(func, specimens))

        # Give list of fitness results back to ES.
        es.tell(fitness_list)

        # get best parameter, fitness from ES
        es_solution = es.result()
        duration = datetime.datetime.now() - start_time

        history = {
            'best_params': es_solution[0],  # Best historical parameters.
            'best_fitness': es_solution[1],  # Best historical reward.
            'curr_best_fitness':
            es_solution[2],  # Best fitness of current generation.
            'mean_fitness':
            fitness_list.mean(),  # Mean fitness of current generation.
            'std_fitness':
            fitness_list.std()  # Std of fitness of current generation.
        }
        np.savez(os.path.join(RESULTS_DIR, 'controller', dir_name, str(i)),
                 **history)

        print(
            'Gen: {0:}\t| Best fit of gen: {1:.2f}\t| Best fit historical: {2:.2f}\t|'
            ' Mean fit: {3:.2f}\t| Std of fit: {4:.2f}\t| Time: {5:}m {6:}s'.
            format(i, es_solution[2], es_solution[1], fitness_list.mean(),
                   fitness_list.std(),
                   *divmod(int(duration.total_seconds()), 60)))

        if i % args.eval_interval == 0:
            start_time = datetime.datetime.now()
            eval_fitness_list = np.array(
                pool.map(
                    func,
                    np.broadcast_to(es_solution[0], (args.n_workers, ) +
                                    es_solution[0].shape)))
            duration = datetime.datetime.now() - start_time
            print(
                '{0:}-worker average fit of best params after gen {1:}: {2:.2f}. Time: {3:}m {4:}s.'
                .format(args.n_workers, i, eval_fitness_list.mean(),
                        *divmod(int(duration.total_seconds()), 60)))

            np.savez(
                os.path.join(RESULTS_DIR, 'controller', dir_name,
                             str(i) + '_eval'), eval_fitness_list)
Ejemplo n.º 2
0
def main():
    parser = argparse.ArgumentParser(description='RNN')
    parser.add_argument('--batch_size',
                        type=int,
                        default=100,
                        help='Input batch size for training (default=100)')
    parser.add_argument('--n_epochs',
                        type=int,
                        default=20,
                        help='Number of epochs to train (default=20)')
    parser.add_argument('--latent_dim',
                        type=int,
                        default=32,
                        help='Dimension of latent space (default=32)')
    parser.add_argument('--seq_len',
                        type=int,
                        default=1000,
                        help='Length of sequences for learning (default=1000)')
    parser.add_argument('--action_dim',
                        type=int,
                        default=3,
                        help='Dimension of action space (default=3)')
    parser.add_argument('--rnn_hidden_dim',
                        type=int,
                        default=256,
                        help='Dimension of RNN hidden state (default=256)')
    parser.add_argument(
        '--n_gaussians',
        type=int,
        default=5,
        help='Number of gaussians for the Mixture Density Network (default=5)')
    parser.add_argument('--learning_rate',
                        type=float,
                        default=1e-3,
                        help='Learning rate for optimizer (default=1e-3)')
    parser.add_argument('--grad_clip',
                        type=float,
                        default=1.0,
                        help='Gradient clipping value (default=1.0)')
    parser.add_argument('--cuda',
                        action='store_true',
                        default=False,
                        help='enables CUDA training')
    parser.add_argument('--vae_fname', help='VAE model file name')
    parser.add_argument('--train_dir_name',
                        help='Rollouts directory name for training')
    parser.add_argument('--test_dir_name',
                        help='Rollouts directory name for testing')
    parser.add_argument('--log_interval',
                        nargs='?',
                        default='2',
                        type=int,
                        help='After how many epochs to log')
    args = parser.parse_args()

    # TODO: is there a better way to do this?
    if not os.path.exists(
            os.path.join(DATA_DIR, 'rollouts', args.train_dir_name)):
        print("Folder {} does not exist.".format(args.train_dir_name))
        pass
    if not os.path.exists(
            os.path.join(DATA_DIR, 'rollouts', args.test_dir_name)):
        print("Folder {} does not exist.".format(args.test_dir_name))
        pass

    # Read in and prepare the data.
    train_dataset = RolloutDataset(
        path_to_dir=os.path.join(DATA_DIR, 'rollouts', args.train_dir_name),
        size=int(args.train_dir_name.split('_')[-1]))  # TODO: hack. fix?
    train_loader = DataLoader(train_dataset,
                              batch_size=args.batch_size,
                              shuffle=True)

    test_dataset = RolloutDataset(
        path_to_dir=os.path.join(DATA_DIR, 'rollouts', args.test_dir_name),
        size=int(args.test_dir_name.split('_')[-1]))  # TODO: hack. fix?
    test_loader = DataLoader(test_dataset,
                             batch_size=args.batch_size,
                             shuffle=True)

    # Use GPU if available.
    use_cuda = args.cuda and torch.cuda.is_available()
    device = torch.device('cuda' if use_cuda else 'cpu')

    # Load the VAE model from file.
    vae = VAE(latent_dim=args.latent_dim)
    vae.load_state_dict(
        torch.load(os.path.join(DATA_DIR, 'vae', args.vae_fname)))
    vae.to(device)

    # Set up the MDNRNN model and the optimizer.
    mdnrnn = MDNRNN(action_dim=args.action_dim,
                    hidden_dim=args.rnn_hidden_dim,
                    latent_dim=args.latent_dim,
                    n_gaussians=args.n_gaussians).to(device)
    optimizer = optim.Adam(params=mdnrnn.parameters(), lr=args.learning_rate)

    # Train procedure.
    def train(epoch):
        mdnrnn.train()
        train_loss = 0
        start_time = datetime.datetime.now()
        for batch_id, batch in enumerate(train_loader):
            obs_batch = batch['obs'].to(device)
            act_batch = batch['act'].to(device)

            optimizer.zero_grad()

            # Encode obs using VAE.
            vae_obs_batch = obs_batch.view(
                (-1, ) + obs_batch.size()[2:])  # Reshape for VAE.
            z_batch = vae.reparameterize(*vae.encode(vae_obs_batch))
            z_batch = z_batch.view(-1, args.seq_len, args.latent_dim)

            # Predict all but first encoded obs from all but last encoded obs and action.
            targets = z_batch[:, 1:]
            z_batch = z_batch[:, :-1]
            act_batch = act_batch[:, :-1]

            pi, mu, sigma, _ = mdnrnn(act_batch, z_batch)

            loss = nll_gmm_loss(targets, pi, mu, sigma)
            loss.backward()
            train_loss += loss.item()

            torch.nn.utils.clip_grad_value_(mdnrnn.parameters(),
                                            args.grad_clip)
            optimizer.step()

            if batch_id % args.log_interval == 0:
                print(
                    'Epoch: {0:}\t| Examples: {1:}/{2:} ({3:.0f}%)\t| Loss: {4:.2f}\t'
                    .format(epoch, (batch_id + 1) * len(obs_batch),
                            len(train_loader.dataset),
                            100. * (batch_id + 1) / len(train_loader),
                            loss.item() / len(obs_batch)))

        duration = datetime.datetime.now() - start_time
        print(
            'Epoch {} average train loss was {:.4f} after {}m{}s of training.'.
            format(epoch, train_loss / len(train_loader.dataset),
                   *divmod(int(duration.total_seconds()), 60)))

    # Test procedure.
    def test(epoch):
        mdnrnn.eval()
        test_loss = 0
        with torch.no_grad():
            for batch_id, batch in enumerate(test_loader):
                obs_batch = batch['obs'].to(device)
                act_batch = batch['act'].to(device)

                # Encode obs using VAE.
                vae_obs_batch = obs_batch.view(
                    (-1, ) + obs_batch.size()[2:])  # Reshape for VAE.
                z_batch = vae.reparameterize(*vae.encode(vae_obs_batch))
                z_batch = z_batch.view(-1, args.seq_len, args.latent_dim)

                # Predict all but first encoded obs from all but last encoded obs and action.
                targets = z_batch[:, 1:]
                z_batch = z_batch[:, :-1]
                act_batch = act_batch[:, :-1]

                pi, mu, sigma, _ = mdnrnn(act_batch, z_batch)

                test_loss += nll_gmm_loss(targets, pi, mu, sigma).item()
            print('Epoch {} average test loss was {:.4f}.'.format(
                epoch, test_loss / len(test_loader.dataset)))

    # Train/test loop.
    for i in range(1, args.n_epochs + 1):
        train(i)
        test(i)

    # Save the learned model.
    if not os.path.exists(os.path.join(DATA_DIR, 'rnn')):
        os.makedirs(os.path.join(DATA_DIR, 'rnn'))

    torch.save(
        mdnrnn.state_dict(),
        os.path.join(
            DATA_DIR, 'rnn',
            datetime.datetime.today().isoformat() + '_' + str(args.n_epochs)))