rlx-swift: RL infrastructure on mlx-swift

[simeetnayan81]

Hi, I’m a junior developer building rlx-swift (https://github.com/simeetnayan81/rlx-swift): RL infrastructure on mlx-swift, with the goal of a practical full RL stack for Apple silicon (environments, collection, training, and standard algorithms over time).

What’s done so far
The first slice is the environment and collection layer—MDP reset / step, spaces, wrappers, explicit seeding, registry, sync/async vector envs, reference envs (e.g. CartPole, Pendulum), contract tests, and DocC:

• RLXCore, RLXWrappers, RLXEnvs, RLXVector, RLXTesting

APIs are oriented around MLXArray, no global RNG in envs, and normal Swift / throws usage.

Where I want this to go
A whole infrastructure, not only envs: policies and optimizers on MLXNN / MLXOptimizers, rollouts and buffers, and algorithms such as PPO, DQN, and similar as the project matures. I started with environments so the interaction surface is solid before training code piles on. Next for me is closing the loop with training and algorithms on top of this layer.

I’d be very happy if rlx-swift could eventually sit under ml-explore or be considered in-scope as a community / satellite project next to mlx-swift, if that ever makes sense for the org. I’m not assuming that, just putting the interest on the table since doing this all alone without any feedback from community/mentors might not lead to fruitful project.

Help I’m looking for
I’m early in my career and building this in the open, so guidance would genuinely help. In particular:

• Direction: Does env-first, then algorithms on MLXNN/MLXOptimizers, sound like a good path, or should we rather focus on first developing the core algorithms and infrasctucture for RL?
• API / design review: Skimming design.md since this is the design I was able to come up with the help of some coding LLMs that best fits what I was aiming to develop
• MLX best practices: Batching, devices/streams (CLI vs app), PRNG usage, and anything that’s easy to get wrong for a newcomer.
• Packaging & quality bar: SPM layout, versioning, CI expectations, docs, and what “healthy” looks like for a library next to mlx-swift.
• Algorithms roadmap: Sensible order (e.g. which algorithm first), and what the env API should expose for rollouts/buffers so we don’t paint ourselves into a corner.
• Mentorship / review: Occasional code review, pairing on hard design choices, or a co-maintainer if anyone is interested.
• Community signal: what the project would need to look like to be considered.
• Benchmarking and positioning ourselves: A good plan for benchmarking rlx-swift against standard libraries for rl in python on apple devices

Repo: https://github.com/simeetnayan81/rlx-swift

Tagging @davidkoski @zcbenz @angeloskath @andresy here.

Thanks for mlx-swift, and for any time or advice you’re able to share.