flexible loss metrics / plot.Loss
christianhidber opened this issue · comments
Loss metric is currently store in TrainContext / ActorCriticContext fields. SacAgent has an additional alpha_loss. Instead of another derived class, we could use a metrics dictionary in TrainContext. Holding the algorithm specific metrics. Passing the metric name to plot.Loss could then plot the chosen metric:
eg plot.Loss('alpha') or multiple losses in the same plot (if reasonable) plot.Loss(['total','critic')