title: How do we print action distributions in RLlib during training?
body: <p>I'm trying to print action distributions at the end of each episode to see what my agent is doing.  I've attempted to do put this is rock_paper_scissors_multiagent.py by including the following method </p>
  
  <pre><code>def on_episode_end(info):
      episode = info["episode"]
      policy = episode._policies['learned']
      print(policy.model.base_model.summary())
  
      obs_space = Tuple((Discrete(3),Discrete(3)))
      prep = get_preprocessor(obs_space)(obs_space)
      curr_state = list((0,1))
      curr_state = tuple(curr_state)
      curr_state = prep.transform(curr_state)
      logits, _ = policy.model.from_batch({"obs": np.array([curr_state])})
      dist = policy.dist_class(logits, policy.model)
      dist.sample()
      print(dist.logp([0]))
  
  </code></pre>
  
  <p>And adding the callback option to tune.run.  However, I get the following error.  Is this how I should be trying to print the policies after each episode?  And if so, what am I doing wrong?  The rock_paper_scissors_multiagent.py original example is here: <a href="https://github.com/ray-project/ray/blob/master/rllib/examples/rock_paper_scissors_multiagent.py" rel="nofollow noreferrer">https://github.com/ray-project/ray/blob/master/rllib/examples/rock_paper_scissors_multiagent.py</a></p>
  
  <pre><code>&gt;Traceback (most recent call last):
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/tune/trial_runner.py", line 515, in _process_trial
      result = self.trial_executor.fetch_result(trial)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/tune/ray_trial_executor.py", line 351, in fetch_result
      result = ray.get(trial_future[0])
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/worker.py", line 2121, in get
      raise value.as_instanceof_cause()
  ray.exceptions.RayTaskError(ValueError): [36mray_worker[39m (pid=5765, host=Charlies-MBP.fios-router.home)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/agents/trainer.py", line 418, in train
      raise e
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/agents/trainer.py", line 407, in train
      result = Trainable.train(self)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/tune/trainable.py", line 176, in train
      result = self._train()
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/agents/trainer_template.py", line 129, in _train
      fetches = self.optimizer.step()
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/optimizers/multi_gpu_optimizer.py", line 140, in step
      self.num_envs_per_worker, self.train_batch_size)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/optimizers/rollout.py", line 29, in collect_samples
      next_sample = ray_get_and_free(fut_sample)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/utils/memory.py", line 33, in ray_get_and_free
      result = ray.get(object_ids)
  ray.exceptions.RayTaskError(ValueError): [36mray_worker[39m (pid=5768, host=Charlies-MBP.fios-router.home)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/evaluation/rollout_worker.py", line 469, in sample
      batches = [self.input_reader.next()]
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/evaluation/sampler.py", line 56, in next
      batches = [self.get_data()]
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/evaluation/sampler.py", line 99, in get_data
      item = next(self.rollout_provider)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/evaluation/sampler.py", line 319, in _env_runner
      soft_horizon, no_done_at_end)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/evaluation/sampler.py", line 473, in _process_observations
      "episode": episode
    File "rock_paper_scissors_multiagent.py", line 204, in on_episode_end
      logits, _ = policy.model.from_batch({"obs": np.array([curr_state])})
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/models/modelv2.py", line 197, in from_batch
      return self.__call__(input_dict, states, train_batch.get("seq_lens"))
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/models/modelv2.py", line 154, in __call__
      res = self.forward(restored, state or [], seq_lens)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/ray/rllib/models/tf/fcnet_v2.py", line 84, in forward
      model_out, self._value_out = self.base_model(input_dict["obs_flat"])
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 634, in __call__
      outputs = call_fn(inputs, *args, **kwargs)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 751, in call
      return self._run_internal_graph(inputs, training=training, mask=mask)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 893, in _run_internal_graph
      output_tensors = layer(computed_tensors, **kwargs)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 586, in __call__
      self.name)
    File "/Users/charliehou/anaconda3/lib/python3.6/site-packages/tensorflow/python/keras/engine/input_spec.py", line 159, in assert_input_compatibility
      ' but received input with shape ' + str(shape))
  ValueError: Input 0 of layer fc_value_1 is incompatible with the layer: expected axis -1 of input shape to have value 3 but received input with shape [1, 6]  
  </code></pre>