# Passing non-tensor data from a custom environment to a model

**URL:** <https://discuss.ray.io/t/passing-non-tensor-data-from-a-custom-environment-to-a-model/727>\
**Category:** RLlib\
**Created:** [February 4, 2021, 4:28pm UTC](https://discuss.ray.io/t/passing-non-tensor-data-from-a-custom-environment-to-a-model/727 "2021-02-04T16:28:49Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![Bam4d](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/bam4d/32/373_2.png) [@Bam4d](https://discuss.ray.io/u/Bam4d)\
**Post date:** [February 4, 2021, 4:28pm UTC](https://discuss.ray.io/t/passing-non-tensor-data-from-a-custom-environment-to-a-model/727/1 "2021-02-04T16:28:49Z")

</div>

Hi there I am the creator of Griddly [Griddly Docs — Griddly 0.3.3 documentation](https://griddly.readthedocs.io/en/latest/) and I am currently trying to write rllib wrappers for these environments.

I am using action masking in combinatorial action spaces which means i need to sample from a number of possible heirarchical actions during the `action_sample_fn`.

The hierarchical action space can be compressed to a dict and then the sampling function goes through the tree to create the appropriate mask for the chosen actions.

My problem is that I can’t seem to find a way to pass this dict structure (which is built during `step()`) to the `action_sample_fn`.

Things I have tried:

1. Pass the dict as a custom space in `input_dict['obs']['action_mask'] = {.....}`  
Preprocessor will only accept tensors because it assumes all data is some sort of tensor
2. Compress the tree to a variable length tensor and pass it as a tensor in `input_dict['obs']['action_mask'] = [....]`  
Preprocessor tries to flatten the array but expects the tensor to be fixed length.

I’ve read that possible I will be able to use the `env.step` `info` variable, but that is not included in the input\_dict when it gets to `action_sample_fn`. (Potentially this is available in 2.0.0.dev)

Are there any ways of doing this? Thanks

---

<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:** [February 8, 2021, 8:29am UTC](https://discuss.ray.io/t/passing-non-tensor-data-from-a-custom-environment-to-a-model/727/2 "2021-02-08T08:29:31Z")

</div>

Hey @Bam4d , thanks for the question. This is indeed handled quite confusingly right now due to the fact that our `action_sample_fn` behaves slightly differently for torch and tf.  
I’m assuming you are on tf and your action\_sampler\_fn receives the flattened “obs” (which includes the action mask) as `obs_batch` argument. You can simply do the following inside your custom action\_sampler\_fn now:

```auto
from ray.rllib.models.modelv2 import restore_original_dimensions
orig_obs = restore_original_dimensions(obs_batch, policy.observation_space, "tf")

```

Now, orig\_obs should be a dict containing the action mask.

---

<div class="post-metadata">

**Author:** ![Bam4d](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/bam4d/32/373_2.png) [@Bam4d](https://discuss.ray.io/u/Bam4d)\
**Post date:** [February 8, 2021, 9:25am UTC](https://discuss.ray.io/t/passing-non-tensor-data-from-a-custom-environment-to-a-model/727/3 "2021-02-08T09:25:49Z")

</div>

Hi Sven thanks for the reply.

I’m using torch and the problem isnt that the data is flattened, its that the data is deleted because its _not_ a fixed tensor.

My action mask is either a variable length tensor, or it needs to be a dict object.

It seems as though its not possible to put non-tensor data (i.e a dict object)into the `input_dict`?

I’ve worked around this issue in a pretty ugly way for now by overriding compute\_actions with a mixin, and then setting `_use_trajectory_view_api` to false, and putting the dict into the `info_batch` instead.

It would be much better if the `info_batch` could just be seen by the action\_sampler\_fn, OR if `input_dict` could store dict data.

---

<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:** [February 8, 2021, 10:58am UTC](https://discuss.ray.io/t/passing-non-tensor-data-from-a-custom-environment-to-a-model/727/4 "2021-02-08T10:58:27Z")

</div>

Ah, got it. Yeah, I’ll create a PR, force-adding the info dicts always, no matter what the view requirements say. Similar to rewards and dones, which are also always part of the train batch, no matter whether they are needed in the loss or not.

---

<div class="post-metadata">

**Author:** ![Bam4d](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/bam4d/32/373_2.png) [@Bam4d](https://discuss.ray.io/u/Bam4d)\
**Post date:** [February 8, 2021, 12:52pm UTC](https://discuss.ray.io/t/passing-non-tensor-data-from-a-custom-environment-to-a-model/727/5 "2021-02-08T12:52:42Z")

</div>

This sounds perfect!

Is there an issue number or anything i can track on github?

Thanks!