diff --git a/config/networks.json5 b/config/networks.json5
index 4471a35..ebce18f 100644
--- a/config/networks.json5
+++ b/config/networks.json5
@@ -6,7 +6,7 @@
         output_dim: 5
     },
     deepsets: {
-        input_dim: 5, // number of relative state vars for others
+        input_dim: 6, // number of relative state vars for others
         phi: {
             hidden_n: 2,
             hidden_dim: 20,
diff --git a/src/bc/bc.py b/src/bc/bc.py
index ed5fe31..418df82 100644
--- a/src/bc/bc.py
+++ b/src/bc/bc.py
@@ -14,7 +14,7 @@ def bc_config(ray_config):
     config = {
         'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0},
         'deepsets': {
-            'input_dim': 5,
+            'input_dim': 6,
             'phi': {
                 'hidden_n': ray_config['deepsets_phi_hidden_n'],
                 'hidden_dim': ray_config['deepsets_phi_hidden_dim']
@@ -39,7 +39,7 @@ def bc_config(ray_config):
             'lr':ray_config['lr'],
             'weight_decay':ray_config['weight_decay']
         },
-        'train_epochs': 50,
+        'train_epochs': 20,
         'train_batch_size': ray_config['train_batch_size'],
         'loss': ray_config['loss'],