# Pytorch Geometric in RLLib?

**URL:** <https://discuss.ray.io/t/pytorch-geometric-in-rllib/3079>\
**Category:** RLlib\
**Created:** [July 30, 2021, 7:27am UTC](https://discuss.ray.io/t/pytorch-geometric-in-rllib/3079 "2021-07-30T07:27:53Z")\
**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:** [July 30, 2021, 7:27am UTC](https://discuss.ray.io/t/pytorch-geometric-in-rllib/3079/1 "2021-07-30T07:27:54Z")

</div>

Hello,

I’m trying to train a Graph Network on RLLib. For this purpose I’m using the Pytorch Geometric library.

This library has a specific way to collate samples together in a batch (because each sample has a graph of different size, more info [here](https://pytorch-geometric.readthedocs.io/en/latest/notes/batching.html#)).

**I’m a bit lost on how to integrate Pytorch Geometric with RLLib.**

* * *

Should I convert observation to tensor graph representation in my custom environment ? Or should I do it on model side ?

Alternatively, can I change the way samples are collated into batch ?

---

<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:** [August 3, 2021, 3:58pm UTC](https://discuss.ray.io/t/pytorch-geometric-in-rllib/3079/2 "2021-08-03T15:58:05Z")

</div>

Hey @Astariul , have you taken a look at the `Repeated` space that RLlib offers? It allows for variable sized observations. There is an example script here that shows how you can use it with an env that outputs variable sized observations:

> <https://github.com/ray-project/ray/blob/master/rllib/examples/complex_struct_space.py>

---

<div class="post-metadata">

**Author:** ![smorad](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/smorad/32/272_2.png) [@smorad](https://discuss.ray.io/u/smorad)\
**Post date:** [August 9, 2021, 12:27pm UTC](https://discuss.ray.io/t/pytorch-geometric-in-rllib/3079/3 "2021-08-09T12:27:17Z")

</div>

> [@Astariul](#):
>
> This library has a specific way to collate samples together in a batch (because each sample has a graph of different size, more info [here](https://pytorch-geometric.readthedocs.io/en/latest/notes/batching.html#)).

I use torch geometric in ray rllib here: [https://github.com/smorad/graph-conv-memory-paper/blob/master/src/models/ray\_graph.py](https://github.com/smorad/graph-conv-memory-paper/blob/master/src/models/ray_graph.py)