Neuromodulation on Weight Manifolds

*CSHL (Zador Lab)

Python, JAX, Flax, Brax, PPO*

Problem: Standard Reinforcement Learning (RL) struggles to generalize to related tasks without extensive retraining.
Method: Implemented a generalized model of neuromodulation that navigates a “weight manifold” rather than optimizing single weights. Engineered custom wrappers in JAX/Brax to integrate this manifold approach with Proximal Policy Optimization (PPO).
Outcome: Created transformation-safe architectures for continuous control tasks and visualized the optimization landscape using Manim. (See Preprint).