How jhunter533/sansac works
1. What this repository is
- This is a reinforcement learning training project centered on Soft Actor-Critic (SAC), with an optional spiking actor variant using SpikingJelly.
- The architecture is small and code-centric: most behavior is driven by a few Python modules rather than a large package hierarchy.
2. Main execution path
- runner.py is the likely top-level entry point for launching training runs.
- The README also shows direct programmatic use, where an Agent-like abstraction is instantiated, trained, evaluated, and saved.
- From an architecture perspective, the runtime flow is: create gymnasium environment → construct agent/manager → train policy/value networks → log metrics → save model → evaluate.
3. Core learning logic
- algorithm.py contains the actual SAC/SANSAC implementation.
- This is the main model boundary: actor/critic networks, policy updates, value updates, entropy-driven SAC logic, and the optional spiking behavior when SAN=True.
- Because the README mentions SpikingJelly, this module is the place where standard neural layers and spike-based layers likely diverge.
4. Training orchestration
- manager.py appears to be the control layer that coordinates environment interaction, replay collection, optimization steps, evaluation, and checkpointing.
- In a typical architecture like this, the manager sits between the environment and the algorithm, handling episodes, timesteps, batch sampling, and when to invoke policy/critic updates.
- This is the main “application service” boundary, separate from the pure learning math in algorithm.py.
5. Configuration boundary
- config.py likely centralizes run parameters such as environment name, hidden dimensions, seed, batch sizes, learning rates, GPU/device flags, and logging options.
- This keeps experiment settings out of the algorithm code and makes runner.py simpler.
- It also helps explain the repo’s flexibility across gymnasium environments: changing dimensions or environment-specific settings should happen here or in the manager layer.
6. Logging and observability
- logger.py provides experiment logging, with Weights & Biases (wandb) as the external tracking system mentioned in the README.
- Logging is explicitly not core to the algorithm, so it should be treated as a sidecar concern: metrics, losses, rewards, maybe system stats.
- This module forms the observability boundary, separate from training correctness.
7. Environment and simulation boundary
- The primary environment is Gymnasium’s BipedalWalker-v3, but the README notes it can be adapted to other Gymnasium tasks by changing input dimensions.
- This means the repo depends on the Gymnasium API contract for reset/step/observation/action spaces.
- The environment is external to the learning system and acts as the source of trajectories and reward signals.
8. Hardware/runtime boundary
- The project targets PyTorch for tensor computation and training.
- It can use CUDA/cupy for GPU acceleration, and the README notes those can be disabled if needed.
- There is also a mention of pynvml, which suggests some GPU monitoring or device management support, likely in the logger or manager path.
9. Data flow through the system
- Observations come from Gymnasium into the agent.
- The policy in algorithm.py produces actions; those actions are sent back to the environment.
- Transitions and rewards accumulate under the manager’s control, then are used to update the actor and critic networks.
- Metrics flow outward to logger.py/wandb, while learned weights flow to save/load checkpoints.
10. Optional spiking path
- The spiking variant is not a separate application; it is an alternate model path inside the same SAC pipeline.
- The likely architecture difference is confined to the actor network implementation, while the surrounding training loop and environment interaction remain the same.
- This makes the project a good example of a shared RL backbone with a switchable policy representation.
11. Artifacts and outputs
- The main outputs are trained model checkpoints and logged experiment runs.
- Evaluation appears to be performed after training, likely using the learned policy in inference mode without exploration noise.
- There is no indication of a web UI or service layer; the primary artifact is the trained agent itself.
12. Repository shape in one sentence
- Architecturally, this repo is a single-process RL training application with clear separation between algorithm logic (algorithm.py), run orchestration (manager.py, runner.py), configuration (config.py), and observability (logger.py), all built around Gymnasium, PyTorch, and optional SpikingJelly/wandb integrations.
Main entry points: manager.py Manager · algorithm.py SAC core · logger.py Logger · runner.py Runner
Components 16
Core training
Main SAC/SANSAC learning and execution flow.
- Runner · Top-level launch script for training runs.
runner.py - Manager · Coordinates rollout collection, updates, evaluation, and checkpointing.
manager.py - SAC core · Implements SAC and the optional spiking actor variant.
algorithm.py - Actor · Produces actions from observations, with optional spiking implementation.
algorithm.py - Critic · Estimates action values for SAC updates.
algorithm.py - Replay buffer · Holds collected transitions for off-policy updates.
manager.py - Checkpoint · Persisted policy and training state artifacts.
manager.py
Supporting runtime
Configuration, logging, and external dependencies around training.
- Config · Centralizes hyperparameters, device settings, and environment options.
config.py - Logger · Handles experiment tracking and reporting to external logging backends.
logger.py - Gymnasium · Source of observations, rewards, and terminal signals.
- BipedalWalker · Primary training task referenced by the project.
- PyTorch · Provides model execution, gradients, and optimization.
- SpikingJelly · Optional spiking neural network components used by the actor variant.
- Weights & Biases · External experiment tracking and visualization service.
- CuPy · Optional GPU-accelerated array/runtime support.
- NVML · Optional GPU/device monitoring support.
Connections 16
- Runner → Config: loads run settings · Runner reads experiment parameters before starting training.
- Runner → Manager: starts · Creates the orchestration layer for the training run.
- Manager → Gymnasium: drives env · Controls episode reset and step interaction through the Gymnasium API.
- Gymnasium → BipedalWalker: task · Primary environment used in the default setup.
- Manager → SAC core: updates · Invokes policy and critic optimization using sampled experience.
- Manager → Replay buffer: stores transitions · Collects experience for off-policy learning.
- Replay buffer → SAC core: samples batch · Provides minibatches for SAC/SANSAC training updates.
- SAC core → Actor: trains · Updates the policy network parameters.
- SAC core → Critic: trains · Updates value estimation networks.
- SAC core → SpikingJelly: optional path · Enables the spiking actor variant when SAN is active.
- Manager → Logger: reports metrics · Sends rewards, losses, and other run statistics.
- Logger → Weights & Biases: syncs · Forwards tracked metrics to the external service.
- SAC core → PyTorch: runs on · Uses PyTorch for tensor operations and optimization.
- Manager → Checkpoint: saves · Writes trained weights and training state to disk.
- Manager → CuPy: optional GPU · May use GPU-accelerated array support when enabled.
- Logger → NVML: monitors device · May inspect GPU/runtime status for logging or diagnostics.