Training

Simple training loop: interact with environment and do training step.

class jax_agents.common.training.TrainConfig(env: Any, algorithm: Any, folder: str, timesteps: int, max_episode_len: int = 200, n_steps: int = 1, buffer_size: int = 100000, batch_size: int = 128, seed: int = 1996)

Bases: object

Config to initialize training loop.

Parameters:

• env – environment to train on
• algorithm – rl algorithm to solve the problem
• folder – path to save the results
• timesteps – how long to train
• max_episode_len – when to reset the environment
• n_steps – support to multistep reinforcement learning
• buffer_size – how many transitions to store
• batch_size – used for training
• seed – random seed

batch_size = 128

buffer_size = 100000

max_episode_len = 200

n_steps = 1

seed = 1996

jax_agents.common.training.train(config: jax_agents.common.training.TrainConfig)

Start the training loop with the given configuration.