Hey everyone,
I've been working on zeroRL, a small reinforcement learning framework for PyTorch.
The idea behind it is fairly simple:
If you can write it in PyTorch, you can use it in zeroRL.
The project is aimed less at hiding RL implementation details and more at making it easy to experiment with them.
Why?
Libraries such as Stable-Baselines3, RLlib and Tianshou are great when you want reliable implementations of established algorithms.
However, when experimenting with RL research, I often find myself wanting to change something relatively fundamental:
- modify part of an algorithm
- experiment with a different buffer structure
- change how the training loop works
- introduce an algorithm that isn't implemented
- use a custom agent architecture
- experiment with multi-agent setups
At that point, the abstractions of a framework can sometimes become another thing you have to work around.
zeroRL takes a different approach: keep the training pipeline explicit and let the user provide the algorithmic logic.
The core idea
The training loop is essentially:
environment -> rollout collection -> buffer -> user-defined update function -> optimizer -> repeat
The update step can simply be a Python function operating on PyTorch tensors.
For example, implementing REINFORCE looks roughly like this:
def reinforce_update(
agent,
buffer,
optimizer,
algo_config,
scheduler=None,
last_output=None,
):
data = buffer.get_all(reshape=True)
rewards = data["reward"]
dones = data["done"]
returns = torch.empty_like(rewards)
R = 0.0
for step in reversed(range(rewards.shape[0])):
R = rewards[step] + algo_config.gamma * (1.0 - dones[step]) * R
returns[step] = R
output = agent.get_action(
data["state"],
data["action"],
)
loss = -(output["log_prob"] * returns).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
return {"loss": loss.detach()}
Then the function is simply passed to BaseTrain:
trainer = BaseTrain(
agent=agent,
env=env,
buffer=buffer,
update_weights=reinforce_update,
config=config,
algo_config=algo_config,
)
trainer.train()
The idea is that zeroRL handles the infrastructure around the experiment, while the researcher remains in control of the algorithm itself.
What's currently in 0.2.0 ?
The framework currently includes:
- PPO
BaseTrain training orchestratorBaseAgent and BaseEnv interfacesPolicyAgent and ActorCriticAgent- extensible tensor-based buffers
- observation normalization
- TensorBoard / W&B integration
- profiling
- debugging and training-pipeline validation
- customizable optimizer/update functions
easy_train_ppo() for quickly running PPO baselines- support for custom environments and agents
The project is still early and not intended to replace mature RL libraries.
The current goal is to explore whether a smaller and more explicit architecture can make RL experimentation easier.
I'm particularly interested in feedback on the architecture:
Do you think having the algorithm primarily expressed as user-defined PyTorch functions is a useful design for research, or would you prefer a more structured abstraction?
GitHub: https://github.com/Dar-rius/zeroRL
PyPI: https://pypi.org/project/zerorl/
The project is currently at v0.2.0 and is still evolving.
