# Distribution wrapper with dict space

**URL:** <https://discuss.ray.io/t/distribution-wrapper-with-dict-space/2292>\
**Category:** RLlib\
**Created:** [May 26, 2021, 9:09am UTC](https://discuss.ray.io/t/distribution-wrapper-with-dict-space/2292 "2021-05-26T09:09:55Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![Astariul](https://avatars.discourse-cdn.com/v4/letter/a/2acd7d/32.png) [@Astariul](https://discuss.ray.io/u/Astariul)\
**Post date:** [May 26, 2021, 9:09am UTC](https://discuss.ray.io/t/distribution-wrapper-with-dict-space/2292/1 "2021-05-26T09:09:55Z")

</div>

My environment have a `spaces.Dict` for the action spaces.

I wrote a custom `DistributionWrapper`, returning a `dict`, and it looks all fine :

```python
class CustomDistribution(TorchDistributionWrapper):
   def sample(self):
        idx_dist = self._idx_distribution()
        idx = idx_dist.sample()

        type_dists = self._type_distribution(idx)
        actions = [d.sample() for d in type_dists]
        self._action_logp = idx_dist.logp(idx) + sum(d.logp(a) for d, a in zip(type_dists, actions))

        return {
            "idx": idx,
            "p1": actions[0],
            "p2": actions[1],
            "p3": actions[2],
        }

```

But I’m having a problem for the implementation of `logp()` method… This method takes actions as input, and it’s a Tensor format, not a dictionary. The tensor has shape `[batch, 4]` (which is normal, I have 4 actions).

* * *

**Does someone know where I can find the code `dict` → `tensor` ?**

Because I’m having the following error :

> ValueError: The value argument must be within the support

Which means the order of my action is messed up !

---

<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 26, 2021, 11:19am UTC](https://discuss.ray.io/t/distribution-wrapper-with-dict-space/2292/2 "2021-05-26T11:19:25Z")

</div>

I think gym sorts Dict actions alphabetically by key (OrderedDict). Could that be the reason?

---

<div class="post-metadata">

**Author:** ![Astariul](https://avatars.discourse-cdn.com/v4/letter/a/2acd7d/32.png) [@Astariul](https://discuss.ray.io/u/Astariul)\
**Post date:** [May 26, 2021, 11:46pm UTC](https://discuss.ray.io/t/distribution-wrapper-with-dict-space/2292/3 "2021-05-26T23:46:20Z")

</div>

Thank you for the answer !  
I checked and indeed the keys seems to be reordered alphabetically.

Is there any way to get a dict of Tensor instead of just a tensor ?