Skip to content

Project Architecture

Kirill Hitushkin edited this page Jul 23, 2026 · 1 revision

Project architecture

Execution flow

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.

Package map

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.

Configured registries

Environment backends, agents, policies, and network nodes are independent extension points built around the same pattern:

  1. a small config contains name and params
  2. a decorator registers a class under name
  3. a builder imports known component modules
  4. the shared factory converts params into the component's typed Config
  5. 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

Extension boundaries

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.

Node-based networks

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.

Explicit network state

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.

Clone this wiki locally