vwxyzjn / enn-trainer

Entity Neural Network Trainer

ENN Trainer allow you to train reinforcement learning agents for Entity Gym environments with PPO or behavioral cloning.

Installation

pip install enn-trainer
pip install setuptools==59.5.0
pip install torch==1.10.2+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
pip install torch-scatter -f https://data.pyg.org/whl/torch-1.10.0+cu113.html

Usage

Training policy for an entity-gym example environment:

python -m  enn_trainer.train env.id=Xor

List all available hyperparameters:

python -m enn_trainer.train --hps-info

Setting up a training script for a custom entity-gym environment (replace TreasureHunt with your environment):

from enn_trainer import TrainConfig, State, init_train_state, train
from entity_gym.examples.tutorial import TreasureHunt
import hyperstate

@hyperstate.stateful_command(TrainConfig, State, init_train_state)
def main(state_manager: hyperstate.StateManager) -> None:
    train(state_manager=state_manager, env=TreasureHunt)

if __name__ == "__main__":
    main()

You can find more detailed guides and an API reference on the documentation website. If you run into issues or have a question, feel free to open an issue or ask on discord.