feat: intermediate work

stoix/base_types.py

@@ -116,6 +116,12 @@ class CoreLearnerState(NamedTuple):
     timestep: TimeStep
 
 
+class CoreOffPolicyLearnerState(NamedTuple):
+    params: Parameters
+    opt_states: OptStates
+    key: chex.PRNGKey
+
+
 class OnPolicyLearnerState(NamedTuple):
     """State of the learner. Used for on-policy learners."""
 

stoix/configs/arch/anakin.yaml

@@ -3,8 +3,8 @@ architecture_name: anakin
 # --- Training ---
 seed: 42  # RNG seed.
 update_batch_size: 1 # Number of vectorised gradient updates per device.
-total_num_envs: 1024  # Total Number of vectorised environments across all devices and batched_updates. Needs to be divisible by n_devices*update_batch_size.
-total_timesteps: 1e7 # Set the total environment steps.
+total_num_envs: 1  # Total Number of vectorised environments across all devices and batched_updates. Needs to be divisible by n_devices*update_batch_size.
+total_timesteps: 1e4 # Set the total environment steps.
 # If unspecified, it's derived from num_updates; otherwise, num_updates adjusts based on this value.
 num_updates: ~ # Number of updates
 

stoix/configs/default/sebulba/default_ff_dqn.yaml

@@ -0,0 +1,11 @@
+defaults:
+  - logger: base_logger
+  - arch: sebulba
+  - system: ff_dqn
+  - network: mlp_dqn
+  - env: envpool/cartpole
+  - _self_
+
+hydra:
+  searchpath:
+    - file://stoix/configs

stoix/systems/q_learning/ff_c51.py → stoix/systems/q_learning/anakin/ff_c51.py

@@ -560,7 +560,7 @@ def run_experiment(_config: DictConfig) -> float:
 
 
 @hydra.main(
-    config_path="../../configs/default/anakin",
+    config_path="../../../configs/default/anakin",
     config_name="default_ff_c51.yaml",
     version_base="1.2",
 )

stoix/systems/q_learning/ff_ddqn.py → stoix/systems/q_learning/anakin/ff_ddqn.py

@@ -547,7 +547,7 @@ def run_experiment(_config: DictConfig) -> float:
 
 
 @hydra.main(
-    config_path="../../configs/default/anakin",
+    config_path="../../../configs/default/anakin",
     config_name="default_ff_ddqn.yaml",
     version_base="1.2",
 )

stoix/systems/q_learning/ff_dqn.py → stoix/systems/q_learning/anakin/ff_dqn.py

@@ -553,7 +553,7 @@ def run_experiment(_config: DictConfig) -> float:
 
 
 @hydra.main(
-    config_path="../../configs/default/anakin",
+    config_path="../../../configs/default/anakin",
     config_name="default_ff_dqn.yaml",
     version_base="1.2",
 )

stoix/systems/q_learning/ff_dqn_reg.py → stoix/systems/q_learning/anakin/ff_dqn_reg.py

@@ -550,7 +550,7 @@ def run_experiment(_config: DictConfig) -> float:
 
 
 @hydra.main(
-    config_path="../../configs/default/anakin",
+    config_path="../../../configs/default/anakin",
     config_name="default_ff_dqn_reg.yaml",
     version_base="1.2",
 )

stoix/systems/q_learning/ff_mdqn.py → stoix/systems/q_learning/anakin/ff_mdqn.py

@@ -550,7 +550,7 @@ def run_experiment(_config: DictConfig) -> float:
 
 
 @hydra.main(
-    config_path="../../configs/default/anakin",
+    config_path="../../../configs/default/anakin",
     config_name="default_ff_mdqn.yaml",
     version_base="1.2",
 )

stoix/systems/q_learning/ff_qr_dqn.py → stoix/systems/q_learning/anakin/ff_qr_dqn.py

@@ -574,7 +574,7 @@ def run_experiment(_config: DictConfig) -> float:
 
 
 @hydra.main(
-    config_path="../../configs/default/anakin",
+    config_path="../../../configs/default/anakin",
     config_name="default_ff_qr_dqn.yaml",
     version_base="1.2",
 )

stoix/systems/q_learning/ff_rainbow.py → stoix/systems/q_learning/anakin/ff_rainbow.py

@@ -648,7 +648,7 @@ def run_experiment(_config: DictConfig) -> float:
 
 
 @hydra.main(
-    config_path="../../configs/default/anakin",
+    config_path="../../../configs/default/anakin",
     config_name="default_ff_rainbow.yaml",
     version_base="1.2",
 )