benchmarl.experiment.Callback

class Callback[source]

Bases: object

A Callback that can be added to experiments. To create your callback, you can inherit from this class and reimplement just the functions you need.

experiment

the experiment associated to the callback

Type:

Experiment

on_setup()[source]

A callback called atexperiment setup.

on_batch_collected(batch: TensorDictBase)[source]

A callback called at the end of every collection step.

Parameters:

batch (TensorDictBase) – batch of collected data

on_train_step(batch: TensorDictBase, group: str) TensorDictBase[source]

A callback called for every training step.

Parameters:
  • batch (TensorDictBase) – tensordict with the training batch

  • group (str) – group name

Returns:

a new tensordict containing the loss values

Return type:

TensorDictBase

on_train_end(training_td: TensorDictBase, group: str)[source]

A callback called at the end of training.

Parameters:
  • training_td (TensorDictBase) – tensordict containing the loss values

  • group (str) – group name

on_evaluation_end(rollouts: List[TensorDictBase])[source]

A callback called at the end of every training step.

Parameters:

rollouts (list of TensorDictBase) – tensordict containing the loss values