From 7b2ca6edc7f0b2a9496f8f1d9f1e6bff4918723c Mon Sep 17 00:00:00 2001
From: Arec <jamgochian95@gmail.com>
Date: Mon, 26 Jul 2021 02:38:50 -0700
Subject: [PATCH] adding tensorboard writer for training loss and cv loss

---
 README.md        |  4 ++++
 requirements.txt |  3 ++-
 src/bc/bc.py     | 16 +++++++++++-----
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index f671af7..da3af1c 100644
--- a/README.md
+++ b/README.md
@@ -33,6 +33,10 @@ You can then train a default behavior cloning policy with the following. Be sure
 ```
 python src/main.py --train
 ```
+You can run tensorboard by running the following and opening `localhost:6006` (or alternatively port-forwarding 6006 from the remote server)
+```
+tensorboard --logdir output/
+```
 You can then test the learned policy with the following, and see the animation file in `output/`:
 ```
 python src/main.py --test
diff --git a/requirements.txt b/requirements.txt
index 3b395b1..8045763 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,4 +3,5 @@ torch
 sklearn
 pytest
 json5
-tqdm
\ No newline at end of file
+tqdm
+tensorboard
diff --git a/src/bc/bc.py b/src/bc/bc.py
index 277ad59..b299c86 100644
--- a/src/bc/bc.py
+++ b/src/bc/bc.py
@@ -2,11 +2,11 @@ import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader
 import pickle
-#from torch.utils.tensorboard import SummaryWriter
+from torch.utils.tensorboard import SummaryWriter
 
 from src.policies import DeepSetsPolicy
 from src.util.transform import MinMaxScaler
-import json5
+from tqdm import tqdm
 
 class BehaviorCloningPolicy():
     """
@@ -115,9 +115,9 @@ def generate_transforms(dataset):
 def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
     
     # hyperparams
-    train_epochs = 100
+    train_epochs = 1000
     cv_every = 10
-    epoch_every = 1
+    epoch_every = 1000
     train_batch_size = 64
     cv_batch_size = 256 # doesn't matter
     learning_rate = 1e-3
@@ -140,7 +140,10 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
     loss_fn = nn.HuberLoss(reduction='sum')
     optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate, weight_decay=weight_decay)
     
-    for i in range(train_epochs):
+    # generate tensorboard writer
+    writer = SummaryWriter(filestr)
+    
+    for i in tqdm(range(train_epochs)):
 
         epoch_loss = 0
         for (batch_idx, batch) in enumerate(training_loader):
@@ -160,6 +163,8 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
             epoch_loss += loss.item() / len(train_dataset)
         
         # Write epoch loss
+        
+        writer.add_scalar('training loss',epoch_loss, i)
         if i % epoch_every == 0:
             print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss))
 
@@ -171,6 +176,7 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
                     pred_action = policy(batch) 
                     loss = loss_fn(pred_action, batch['action'])
                     cv_loss += loss.item() / len(cv_dataset)
+            writer.add_scalar('cv loss', cv_loss, i)
             print('Epoch: {}, CV Loss: {}'.format(i, cv_loss))
 
     policy.save_model(filestr)