Addressing comments

.github/workflows/run-debug.yaml

@@ -28,4 +28,4 @@ jobs:
         pip install --upgrade "jax[cpu]==0.4.13"
 
     - name: Train on debug dataset
-      run: WANDB_MODE=disabled python train.py --config experiments/dibya/debug_config.py --name debug
+      run: WANDB_MODE=disabled python train.py --config tests/debug_config.py --name debug

README.md

@@ -31,7 +31,7 @@ See the [Jax Github page](https://github.com/google/jax) for more details on ins
 
 Test the installation by training on the debug dataset:
 ```
-python train.py --config experiments/dibya/debug_config.py  --name debug
+python train.py --config tests/debug_config.py --debug
 ```
 
 ## Training
@@ -67,7 +67,7 @@ Steps to contribute:
 1. Fork the repo and create your branch from `master`.
 2. Use `pre-commit` to enable code checks and auto-formatting.
 3. Test that a basic training starts with the debug dataset with: ```
-python train.py --config experiments/dibya/debug_config.py  --name debug
+python train.py --config tests/debug_config.py
 ```
 
 

config.py

@@ -12,7 +12,13 @@ def update_config(config, **kwargs):
 
 
 def wrap(f):
-    """Simple wrapper to enable passing config strings to `get_config`"""
+    """Simple wrapper to enable passing config strings to `get_config`
+
+    Usage:
+
+    python train.py --config=config.py:vit_s,multimodal
+    python train.py --config=config.py:transformer_size=vit_s
+    """
 
     def wrapped_f(config_string=None):
         if config_string is None:

experiments/dibya/debug_config.py → tests/debug_config.py

@@ -9,7 +9,7 @@ def get_config():
     del base_config["dataset_kwargs"]["oxe_kwargs"]
     config = update_config(
         base_config,
-        num_steps=20,
+        num_steps=2,
         optimizer=dict(
             learning_rate=dict(
                 warmup_steps=1,
@@ -18,16 +18,16 @@ def get_config():
         batch_size=64,
         shuffle_buffer_size=1000,
         num_val_batches=1,
-        log_interval=10,
-        eval_interval=10,
+        log_interval=1,
+        eval_interval=2,
         eval_datasets=None,
         trajs_for_metrics=1,
         trajs_for_viz=1,
         dataset_kwargs={
             "data_kwargs_list": [
                 {
                     "name": "bridge_dataset",
-                    "data_dir": "./datasets/debug_dataset",
+                    "data_dir": "./tests/debug_dataset",
                     "image_obs_keys": ["image_0"],
                     "state_obs_keys": ["state"],
                 },