-
Notifications
You must be signed in to change notification settings - Fork 0
Project Architecture
The spiking-rl-lab command enters spiking_rl_lab.run:main. Hydra composes a typed BaseConfig, then Runner owns the experiment lifecycle:
Hydra config
↓
Runner
├── seed PyTorch/skrl
├── configure MLflow/DagsHub
├── build environment backend → Gymnasium → skrl wrapper
├── build agent → policy network + action policy + optimizer
└── build skrl trainer → train or evaluate
The runner closes the environment even if component or trainer construction fails. During training it records reproducibility artifacts and evaluates the best checkpoint when one exists.
Application code follows a src layout under src/spiking_rl_lab/:
| Path | Responsibility |
|---|---|
run.py |
CLI and Hydra entry point |
app/config.py |
Typed root, runner, trainer, and Optuna configuration |
app/runner.py |
Training/evaluation lifecycle and checkpoint loading |
app/tracking.py |
DagsHub/MLflow setup and reproducibility artifacts |
core/ |
Registry factories, validation, and project exceptions |
envs/ |
Environment backend contract, registry, builder, and implementations |
agents/ |
Learning-agent contract, registry, builder, and algorithm implementations |
policies/ |
Action-distribution contract, registry, builder, and adapters |
networks/ |
Extensible node-network assembly, node registry, tensor shapes, and explicit state utilities |
configs/ |
Hydra groups and retained experiment configurations |
Generated runs belong in runs/. The fallback MLflow database belongs in experiments/. Neither is source code.
Environment backends, agents, policies, and network nodes are independent extension points built around the same pattern:
- a small config contains
nameandparams - a decorator registers a class under
name - a builder imports known component modules
- the shared factory converts
paramsinto the component's typedConfig - the class is constructed with explicitly supplied dependencies.
For example:
policy:
name: categorical
params: {}selects the class registered with @register_policy("categorical"). Invalid names, parameters, inheritance, or dependencies are reported as component-specific creation errors. A new implementation normally requires a typed config, a registered class, inclusion of its module in the matching builder, and an optional Hydra preset. The runner does not need a branch for every concrete algorithm, environment, policy, or node.
The following table is a snapshot of registrations available now, not a closed list of supported component types:
| Component | Registered names |
|---|---|
| Environment backend | gymnasium |
| Agent | reinforce |
| Policy |
categorical, gaussian, deterministic
|
| Network node |
linear, conv1d, conv2d, torch_activation, lif, li
|
Each component family has a narrow responsibility:
- an environment backend constructs and wraps environments without defining the learning algorithm
- an agent owns interaction-time state, rollout or replay storage, optimization, and algorithm-specific metrics
- a policy converts network features into an action distribution without defining the network architecture
- a network node transforms tensors and, when needed, exposes explicit state initialization and reset behavior
- Hydra presets select and compose registered components without making them hard-coded defaults in the runner.
This separation allows new algorithms to reuse environments, policies, and nodes, and allows new environments or neuron models to be introduced without modifying existing algorithms when their contracts are compatible. See Development for the registration workflow.
NodeNetwork builds an ordered ModuleList from YAML. Every node receives the previous node's semantic TensorShape and exposes its own output shape, so incompatible compositions fail during construction rather than training.
Supported shapes omit the batch dimension:
-
DenseTensorShape:[features] -
SequenceTensorShape:[channels, length] -
ImageTensorShape:[channels, height, width].
The current REINFORCE implementation flattens observations before the policy network, so its configured network starts with a dense shape. This is an implementation choice of that agent rather than a limitation of NodeNetwork: other agents can provide sequence or image shapes and use convolution nodes directly.
All node calls follow a functional interface:
outputs, next_state = node(inputs, state)Stateless nodes return None. LIF and LI nodes return Norse state objects. NodeNetwork keeps one state entry per node and does not mutate the input state. An agent that uses a stateful network is responsible for the live state. The current REINFORCE integration:
- initialize it for the observation batch
- advance it on each environment step
- reset rows whose environments terminated or were truncated
- save the state at the beginning of a rollout
- detach it at truncated-BPTT boundaries during replay.
The explicit interface lets different algorithms compose artificial and spiking nodes without hiding temporal state in module buffers. Algorithm-specific state handling is documented on each page under Algorithms.