In classification runs, at the beginning a circuit diagram is plotted (and logged to wandb):
In the RL framework, this was not possible earlier as the analysis required classification specific input.
Might be worth checking if there's an easy way to implement this now and if not see if the analysis require further refactoring.