How can I access tune training stats tensors in a callback?

In tensorboard I could see tune/episodes_total, tune/episodes_this_iter` and many other tensor values related to training. I want to override the value of some of those, and add new ones to track. How can I access those in a callback object supplied via trainer config (since this seems to be the only way to interact with training)? I tried the following which didn’t work:

  1. in on_train_results, access the tf session of the current policy which ended up with an assert
  2. use tune.get_session which didn’t crash but none of the tensors I added showed up.

Hey @bzs, thanks for asking this question. Could you specify in more detail, what exactly you are trying to access using the policy’s session inside on_train_result?

I think we should try using on_learn_on_batch(self, *, policy, train_batch), but at the same time fix it by adding the TFRunBuilder to its args. This way, you can use the builder to add data to the feed-dict and tf-op-fetches (which will then show up in the metrics results).
Let me know, what you think. I’ll make that change right now and update our example script (rllib/examples/custom_metrics_and_callbacks.py).

Hey @sven1977, I’m trying to access the 2 tensor in this screenshot

Is on_learn_on_batch called for both the offline dataset codepath (using InputReader) and the online training path (using a regular Env)? If yes then what you proposed sounds exactly like what I hoped to get :slight_smile: