# \[rllib\] Retrieve and modify the computed discrete action logits to PPO agent

**URL:** <https://discuss.ray.io/t/rllib-retrieve-and-modify-the-computed-discrete-action-logits-to-ppo-agent/2018>\
**Category:** RLlib\
**Created:** [May 3, 2021, 1:43am UTC](https://discuss.ray.io/t/rllib-retrieve-and-modify-the-computed-discrete-action-logits-to-ppo-agent/2018 "2021-05-03T01:43:10Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![lucas\_spangher](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/lucas_spangher/32/938_2.png) [@lucas\_spangher](https://discuss.ray.io/u/lucas_spangher)\
**Post date:** [May 3, 2021, 1:43am UTC](https://discuss.ray.io/t/rllib-retrieve-and-modify-the-computed-discrete-action-logits-to-ppo-agent/2018/1 "2021-05-03T01:43:11Z")

</div>

I’d like to retrieve the logits that the PPO agent’s policy computes for each multidiscrete action, and then modify them so that the action sampling can use my modified logits.

I’m really sorry for the question, but the code is so abstracted that it’s hard for me to tell how I would retrieve the logits and then reset them. Is it in the forward pass of ModelV2? Any direction you could provide would be helpful, thank you.

---

<div class="post-metadata">

**Author:** ![sven1977](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/sven1977/32/53_2.png) [@sven1977](https://discuss.ray.io/u/sven1977)\
**Post date:** [May 4, 2021, 8:22am UTC](https://discuss.ray.io/t/rllib-retrieve-and-modify-the-computed-discrete-action-logits-to-ppo-agent/2018/2 "2021-05-04T08:22:18Z")

</div>

Hey @lucas_spangher , thanks for the question. You are right, it’s not super intuitive how to get there. Action logits are processed inside the Policy.exploration object. PPO uses the StochasticSampling exploration module under: `ray.rllib.utils.exploration.stochastic_sampling.py`

What you should do, it you want to modify these before the sampling step is to sub-class `StochasticSampling`, then override its `get_exploration_action` method applying your logic on the logits modification and make sure it does the sampling on those modified logits, then returns the actions (just like the parent StochasticSampling does it).

Then in your config, just do:

```auto
config:
    exploration_config:
        type: [full path to your new class, e.g. "my_dir.my_exploration.MyExploration"]
        [other c'tor args for your class]

```

This would tell PPO to use your own exploration class, instead.

---

<div class="post-metadata">

**Author:** ![sven1977](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/sven1977/32/53_2.png) [@sven1977](https://discuss.ray.io/u/sven1977)\
**Post date:** [May 4, 2021, 8:23am UTC](https://discuss.ray.io/t/rllib-retrieve-and-modify-the-computed-discrete-action-logits-to-ppo-agent/2018/3 "2021-05-04T08:23:44Z")

</div>

Simplifying our “hooks” into the RLlib code is on our list of things to fix in the short term. This includes providing better (and more flexible) callback options to users via an event-registry-based system, as well as using this system to obsolete e.g. the exploration API and other currently hard-coded hooks.

---

<div class="post-metadata">

**Author:** ![lucas\_spangher](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/lucas_spangher/32/938_2.png) [@lucas\_spangher](https://discuss.ray.io/u/lucas_spangher)\
**Post date:** [May 5, 2021, 1:45am UTC](https://discuss.ray.io/t/rllib-retrieve-and-modify-the-computed-discrete-action-logits-to-ppo-agent/2018/4 "2021-05-05T01:45:04Z")

</div>

Thanks for the response! I’m going to give it a try. Here what I’m trying to do (implement an ordinal response) so if there’s interest in incorporating that I’d be happy to submit a PR:

> **[Discretizing Continuous Action Space for On-Policy Optimization](https://arxiv.org/abs/1901.10500)**
>
> In this work, we show that discretizing action space for continuous control
> is a simple yet powerful technique for on-policy optimization. The explosion in
> the number of discrete actions can be efficiently addressed by a policy with
> factorized...

---

<div class="post-metadata">

**Author:** ![lucas\_spangher](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/lucas_spangher/32/938_2.png) [@lucas\_spangher](https://discuss.ray.io/u/lucas_spangher)\
**Post date:** [May 5, 2021, 2:27am UTC](https://discuss.ray.io/t/rllib-retrieve-and-modify-the-computed-discrete-action-logits-to-ppo-agent/2018/5 "2021-05-05T02:27:37Z")

</div>

Just to confirm, the action\_dist is the logits, correct?

---

<div class="post-metadata">

**Author:** ![mannyv](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/mannyv/32/606_2.png) [@mannyv](https://discuss.ray.io/u/mannyv)\
**Post date:** [May 5, 2021, 2:45am UTC](https://discuss.ray.io/t/rllib-retrieve-and-modify-the-computed-discrete-action-logits-to-ppo-agent/2018/6 "2021-05-05T02:45:47Z")

</div>

Hi @lucas_spangher,

The action distribution depends on your action space. It will be one of these: [ray/tf\_action\_dist.py at master · ray-project/ray · GitHub](https://github.com/ray-project/ray/blob/master/rllib/models/tf/tf_action_dist.py) or the torch equivalent.

If it were Categorical for example, that distribution has an inputs member variable that would be the logits. You could modify those and reassign.`action_dist.inputs = action_dist.inputs + 0`. I think that would be all you need in the constructor of your subclass of StochasticSampling.

Your other option is to implement your own custom model and then modify the logits as your last operation before returning them from the forward method. As I think about it more that might be a more straightforward approach if you want to support multiple different kinds of action distributions.

---

<div class="post-metadata">

**Author:** ![lucas\_spangher](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/lucas_spangher/32/938_2.png) [@lucas\_spangher](https://discuss.ray.io/u/lucas_spangher)\
**Post date:** [May 5, 2021, 4:36am UTC](https://discuss.ray.io/t/rllib-retrieve-and-modify-the-computed-discrete-action-logits-to-ppo-agent/2018/7 "2021-05-05T04:36:17Z")

</div>

Hey mannyv, thank you for the answer. This was incredibly clear and straightforward. I should have no trouble implementing what I want to implement. I will keep you posted when I’m able to get it.
