getting training and testing loop working, adding tqdm to simulator, and reduced number of frames, updating readme

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])

requirements.txt

@@ -3,3 +3,4 @@ torch
 sklearn
 pytest
 json5
+tqdm

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

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: 
+    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)
+            # 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)