Learner takes the arguments

• env: RL environment
• model: A Policy model
• device: Per default picks mps or cuda if available else cpu
• dtype: Per default torch.bfloat16 if device is cuda else torch.float32
• compile_no_lstm: Speedup via torch.compile if model has no lstm
• **kwargs: Passed to the Policy, e.g. hidden_size or lstm

Learner.fit takes the arguments

• iters: Number of iterations
• steps: Number of steps in rollout
• desc: Progress bar description (e.g. 'reward')
• log: If True, tensorboard logging is enabled
  • run tensorboard --logdir=runsand visit http://localhost:6006 in the browser!
• stop_func: Function that stops training if it returns True e.g.

...
def stop(kl, **kwargs):
  return kl > .1

curves = learn.fit(40, steps=16, stop_func=stop)
...

• lr, anneal_lr & args of ppo after bs: Hyperparameters

The most important functions in flashrl/utils.py are

• print_curve: Visualizes the loss across the iters
• play: Plays the environment in the terminal and takes
  • model: A Policy model
  • playable: If True, allows you to act (or decide to let the model act)
  • steps: Number of steps
  • fps: Frames per second
  • obs: Argument of the env that should be rendered as observations
  • dump: If True, no frame refresh -> Frames accumulate in the terminal
  • idx: Agent index between 0 and n_agents (default: 0)