diff --git a/README.md b/README.md
index e46ac83..8d79681 100644
--- a/README.md
+++ b/README.md
@@ -83,7 +83,7 @@ python scripts/b_create_env.py -b building_blocks/enamine_catalog.smi -o envs/ca
 
 #### Option B: Download Prepared Files
 ```bash
-cd experiments/data
+cd experiments/data/envs
 gdown https://drive.google.com/uc?id=192RuXBzM51Mk__-kSKcCs4kthnKHWAgm
 tar -xzvf stock.tar.gz
 ```
@@ -122,6 +122,7 @@ wandb agent <sweep-id>
 1. Get `crossdocked.tar.gz` from [here](https://drive.google.com/file/d/1BKYx_H1m-TzG_75Gk-7sjPkt5ow-Acdw/view?usp=sharing)
 2. Extract dataset:
 ```bash
+cd experiments/data/
 gdown 1BKYx_H1m-TzG_75Gk-7sjPkt5ow-Acdw
 tar -xzvf crossdocked.tar.gz
 ```
diff --git a/experiments/scripts/exp3A_sbdd_proxy.py b/experiments/scripts/exp3A_sbdd_proxy.py
index 9edf12b..1930dfb 100644
--- a/experiments/scripts/exp3A_sbdd_proxy.py
+++ b/experiments/scripts/exp3A_sbdd_proxy.py
@@ -21,8 +21,8 @@ if __name__ == "__main__":
     config.cgflow.use_predicted_pose = True
     config.cgflow.num_inference_steps = 100
 
-    config.task.pocket_conditional.pocket_dir = "./data/experiments/CrossDocked2020/crossdocked_pocket10/"
-    config.task.pocket_conditional.train_key = "./data/experiments/CrossDocked2020/train_keys.csv"
+    config.task.pocket_conditional.pocket_dir = "./data/CrossDocked2020/"
+    config.task.pocket_conditional.train_key = "./data/CrossDocked2020/train_keys.csv"
 
     trainer = Proxy_MultiPocket_Trainer(config)
     trainer.run()
diff --git a/experiments/scripts/exp3B_sbdd_unidock.py b/experiments/scripts/exp3B_sbdd_unidock.py
index 5369418..8c27ed9 100644
--- a/experiments/scripts/exp3B_sbdd_unidock.py
+++ b/experiments/scripts/exp3B_sbdd_unidock.py
@@ -21,8 +21,8 @@ if __name__ == "__main__":
     config.cgflow.use_predicted_pose = True
     config.cgflow.num_inference_steps = 100
 
-    config.task.pocket_conditional.pocket_dir = "./data/experiments/CrossDocked2020/crossdocked_pocket10/"
-    config.task.pocket_conditional.train_key = "./data/experiments/CrossDocked2020/train_keys.csv"
+    config.task.pocket_conditional.pocket_dir = "./data/CrossDocked2020/"
+    config.task.pocket_conditional.train_key = "./data/CrossDocked2020/train_keys.csv"
 
     trainer = UniDock_MultiPocket_Trainer(config)
     trainer.run()
diff --git a/experiments/scripts/exp3C_sbdd_localopt.py b/experiments/scripts/exp3C_sbdd_localopt.py
index fa06ed2..a38a8d6 100644
--- a/experiments/scripts/exp3C_sbdd_localopt.py
+++ b/experiments/scripts/exp3C_sbdd_localopt.py
@@ -21,8 +21,8 @@ if __name__ == "__main__":
     config.cgflow.use_predicted_pose = True
     config.cgflow.num_inference_steps = 100
 
-    config.task.pocket_conditional.pocket_dir = "./data/experiments/CrossDocked2020/crossdocked_pocket10/"
-    config.task.pocket_conditional.train_key = "./data/experiments/CrossDocked2020/train_keys.csv"
+    config.task.pocket_conditional.pocket_dir = "./data/CrossDocked2020/"
+    config.task.pocket_conditional.train_key = "./data/CrossDocked2020/train_keys.csv"
 
     trainer = AutoDock_MultiPocket_Trainer(config)
     trainer.run()
diff --git a/experiments/scripts/exp3Z_sampling.py b/experiments/scripts/exp3Z_sampling.py
index 7b0e3bc..1a09ff9 100644
--- a/experiments/scripts/exp3Z_sampling.py
+++ b/experiments/scripts/exp3Z_sampling.py
@@ -11,7 +11,7 @@ from tqdm import tqdm
 from synthflow.config import Config, init_empty
 from synthflow.pocket_conditional.sampler import PocketConditionalSampler
 
-POCKET_DIR = Path("./data/experiments/CrossDocked2020/crossdocked_pocket10/")
+POCKET_DIR = Path("./data/CrossDocked2020/")
 TEST_KEY_PATH = Path("./data/experiments/CrossDocked2020/test_keys.csv")
 
 
diff --git a/scripts/A_semlaflow_train_ar_plinder.sh b/scripts/A_semlaflow_train_ar_plinder.sh
index 74124f0..2a5448d 100644
--- a/scripts/A_semlaflow_train_ar_plinder.sh
+++ b/scripts/A_semlaflow_train_ar_plinder.sh
@@ -1,26 +1,33 @@
-DATA_DIR="/projects/jlab/to.shen/CGFlow/data/complex/plinder/smol"
+DATA_DIR="experiments/data/complex/plinder_15A"
 
-python scripts/_a1_cgflow_train.py \
+python scripts/_a1_semlaflow_train.py \
   --data_path ${DATA_DIR} \
   --dataset plinder \
   --categorical_strategy auto-regressive \
   --ordering_strategy connected \
   --decomposition_strategy reaction \
   --pocket_n_layers 4 \
+  --n_coord_sets 64 \
   --d_message 64 \
   --d_message_hidden 96 \
   --time_alpha 1.0 \
-  --t_per_ar_action 0.3 \
-  --max_interp_time 0.4 \
-  --max_action_t 0.6 \
+  --t_per_ar_action 0.33 \
+  --max_interp_time 1.0 \
+  --max_action_t 0.66 \
   --max_num_cuts 2 \
   --dist_loss_weight 0. \
   --type_loss_weight 0. \
   --bond_loss_weight 0. \
   --charge_loss_weight 0. \
   --optimal_transport None \
-  --monitor val-strain \
+  --monitor val-conformer-no-align-rmsd \
   --monitor_mode min \
   --val_check_epochs 1 \
-  --batch_cost 600 \
-  --num_gpus 8
+  --acc_batches 2 \
+  --batch_cost 6 \
+  --num_workers 8 \
+  --ligand_local_connections 30 \
+  --pocket_local_connections 30 \
+  --pocket_ligand_local_connections 30 \
+  --max_atoms 800 \
+  --num_gpus 1
\ No newline at end of file
diff --git a/src/cgflow/models/fm.py b/src/cgflow/models/fm.py
index c83f907..31c4534 100644
--- a/src/cgflow/models/fm.py
+++ b/src/cgflow/models/fm.py
@@ -536,26 +536,6 @@ class MolecularCFM(L.LightningModule):
         self.builder = builder
         self.ema_gen = ema_gen if use_ema else None
 
-        # Anything else passed into kwargs will also be saved
-        # hparams = {
-        #     "lr": lr,
-        #     "coord_scale": coord_scale,
-        #     "dist_loss_weight": dist_loss_weight,
-        #     "type_loss_weight": type_loss_weight,
-        #     "bond_loss_weight": bond_loss_weight,
-        #     "type_strategy": type_strategy,
-        #     "bond_strategy": bond_strategy,
-        #     "self_condition": self_condition,
-        #     "distill": distill,
-        #     "lr_schedule": lr_schedule,
-        #     "sampling_strategy": sampling_strategy,
-        #     "use_ema": use_ema,
-        #     "compile_model": compile_model,
-        #     "warm_up_steps": warm_up_steps,
-        #     **gen.hparams,
-        #     **integrator.hparams,
-        #     **kwargs,
-        # }
         self.save_hyperparameters(vars(args))
 
         stability_metrics = {
diff --git a/src/synthflow/tasks/config.py b/src/synthflow/tasks/config.py
index 4c6f155..912931a 100644
--- a/src/synthflow/tasks/config.py
+++ b/src/synthflow/tasks/config.py
@@ -74,9 +74,10 @@ class PocketConditionalConfig:
         Proxy Key from PharmacoNet
     """
 
-    proxy: tuple[str, str, str] = ("TacoGFN_Reward", "QVina", "CrossDocked2020")
-    pocket_dir: str = "./data/experiments/CrossDocked2020/crossdocked_pocket10/"
-    train_key: str = "./data/experiments/CrossDocked2020/train_keys.csv"
+    proxy: tuple[str, str,
+                 str] = ("TacoGFN_Reward", "QVina", "CrossDocked2020")
+    pocket_dir: str = "./data/CrossDocked2020/"
+    train_key: str = "./data/CrossDocked2020/train_keys.csv"
 
 
 @dataclass
@@ -84,4 +85,5 @@ class TasksConfig:
     moo: MOOTaskConfig = field(default_factory=MOOTaskConfig)
     constraint: ConstraintConfig = field(default_factory=ConstraintConfig)
     docking: DockingTaskConfig = field(default_factory=DockingTaskConfig)
-    pocket_conditional: PocketConditionalConfig = field(default_factory=PocketConditionalConfig)
+    pocket_conditional: PocketConditionalConfig = field(
+        default_factory=PocketConditionalConfig)