From 5b09374c772434a10a5ae77cc2c7c8c33c62fd3a Mon Sep 17 00:00:00 2001 From: Arec Date: Fri, 23 Jul 2021 04:44:50 -0700 Subject: [PATCH] getting training and testing loop working, adding tqdm to simulator, and reduced number of frames, updating readme --- README.md | 16 +++++++++++++++- requirements.txt | 1 + src/bc/bc.py | 13 ++++++++----- src/main.py | 34 ++++++++++++++++++++++------------ 4 files changed, 46 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 9116b1e..f671af7 100644 --- a/README.md +++ b/README.md @@ -24,8 +24,22 @@ Once the repository has been set up, you can process and save expert track demon ``` python src/expert_data.py --loc [LOCNUM] --track [TRACKNUM] ``` +You can (and should) process all tracks at once at location 0 with: +``` +python src/expert_data.py --all-tracks +``` -You can then load the experts actions and observations using +You can then train a default behavior cloning policy with the following. Be sure to check help for main.py for running options. +``` +python src/main.py --train +``` +You can then test the learned policy with the following, and see the animation file in `output/`: +``` +python src/main.py --test +``` + + +You can load the experts actions manually ``` from src import expert_data observations, actions = expert_data.load_expert_data(loc = [LOCNUM], track = [TRACKNUM]) diff --git a/requirements.txt b/requirements.txt index 229fb02..3b395b1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ torch sklearn pytest json5 +tqdm \ No newline at end of file diff --git a/src/bc/bc.py b/src/bc/bc.py index 4c0e120..277ad59 100644 --- a/src/bc/bc.py +++ b/src/bc/bc.py @@ -115,8 +115,9 @@ def generate_transforms(dataset): def train(train_dataset, cv_dataset, policy, filestr, **kwargs): # hyperparams - train_epochs = 10000 - cv_every = 100 + train_epochs = 100 + cv_every = 10 + epoch_every = 1 train_batch_size = 64 cv_batch_size = 256 # doesn't matter learning_rate = 1e-3 @@ -138,8 +139,7 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs): # generate loss function, optimizer loss_fn = nn.HuberLoss(reduction='sum') optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate, weight_decay=weight_decay) - pickle.dump(train_dataset[150:160], open(filestr+'_test_batch.pkl', 'wb')) - policy.save_model(filestr) + for i in range(train_epochs): epoch_loss = 0 @@ -147,6 +147,9 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs): # sample mini-batch and run through policy pred_action = policy(batch) + if i == 0 and batch_idx==0: + pickle.dump(batch, open(filestr+'_test_batch.pkl', 'wb')) + policy.save_model(filestr) loss = loss_fn(pred_action, batch['action']) # compute loss and step optimizer @@ -157,7 +160,7 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs): epoch_loss += loss.item() / len(train_dataset) # Write epoch loss - if i % 10 == 0: + if i % epoch_every == 0: print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss)) # measure cv loss diff --git a/src/main.py b/src/main.py index 4cb93b1..e276f62 100644 --- a/src/main.py +++ b/src/main.py @@ -5,9 +5,10 @@ import numpy as np import json5 import os opj = os.path.join - +from tqdm import tqdm from src import InteractionDatasetSingleAgent, metrics +from intersim.utils import get_map_path, get_svt def basestr(**kwargs): """ @@ -60,7 +61,7 @@ def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs # make policy, train and test datasets, and send to policy = policy_class(config) - train_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[0])#,1,2]) + train_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[0,1,2]) cv_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[3]) train_fn(train_dataset, cv_dataset, policy, filestr, **kwargs) @@ -72,14 +73,14 @@ def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs # simulate policy track = 4 - simulate_policy(policy, loc=loc, track=track, filestr=filestr) + simulate_policy(policy, loc=loc, track=track, filestr=filestr, nframes=500) # run test metrics test_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[track]) metrics(filestr, test_dataset, policy) -def simulate_policy(policy, loc=0, track=0, filestr=''): +def simulate_policy(policy, loc=0, track=0, filestr='', nframes=float('inf')): """ Simulate a trained policy Args: @@ -89,20 +90,29 @@ def simulate_policy(policy, loc=0, track=0, filestr=''): filestr (str): path prefix to save simulation to """ # animate from environment - env = gym.make('intersim:intersim-v0', loc=loc, track=track, + svt, svt_path = get_svt(base='InteractionSimulator', loc=loc, track=track) + osm = get_map_path(base='InteractionSimulator', loc=loc) + env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm, min_acc=-np.inf, max_acc=np.inf) + # env = gym.make('intersim:intersim-v0', loc=loc, track=track, + # min_acc=-np.inf, max_acc=np.inf) ob, _ = env.reset() env.render() done = False - while not done: - - # get action - action = policy(ob) + i = 0 + with tqdm(total=min(nframes, env._svt.Tind)) as pbar: + while not done and i < nframes: + i += 1 + + # get action + action = policy(ob) - # propagate environment - ob, r, done, info = env.step(action) - env.render() + # propagate environment + ob, r, done, info = env.step(action) + env.render() + + pbar.update() env.close(filestr=filestr)