# Efficient set and graph space for RL

**URL:** <https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014>\
**Category:** RLlib\
**Created:** [February 11, 2022, 10:30am UTC](https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014 "2022-02-11T10:30:19Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![vakker00](https://avatars.discourse-cdn.com/v4/letter/v/b5e925/32.png) [@vakker00](https://discuss.ray.io/u/vakker00)\
**Post date:** [February 11, 2022, 10:30am UTC](https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014/1 "2022-02-11T10:30:19Z")

</div>

Hi,

I’m working on RL problems that involve dynamically changing input spaces (sets and graphs).  
I’m using the `Repeated` space, as that’s design for this purpose, i.e. to represent a graph (following the PyG convention):

```python
        d = {
           "x": Repeated(
                gym.spaces.Box(-1, 1, shape=x_shape, dtype=np.float32),
                MAX_ELEMS,
            ),
            "edge_index": Repeated(
                gym.spaces.Box(0, MAX_ELEMS, shape=(2,), dtype=np.int64),
                MAX_ELEMS ** 2,
            ),
        }

        observation_space = gym.spaces.Dict(d)

```

It works fine in general. However, it creates a significant overhead, because the allocated tensors scale with `x_shape*MAX_ELEMS + 2*MAX_ELEMS**2`.

I understand that the `flattened` observations need to be fixed, but I never use that anywhere. This creates a large memory and latency overhead on the GPU, because the whole zero-padded tensor is transferred from CPU to GPU. This is especially bad when the number of elements follow a long tailed Poisson distribution, because I need to use a large `MAX_ELEMS` while most of the time there’s only a fraction of the allocated tensor utilized.

I created a graph space for this purpose, such as:

```python
class GraphSpace(gym.Space):
    def __init__ (self, x_shape, e_shape=None, dtype=np.float32):
        super(). __init__ ()
        self.x_shape = x_shape
        self.e_shape = e_shape
        self._shape = (0,)
        self.dtype = dtype

    def sample(self):
        if self.e_shape:
            edge_attr = self.np_random.normal(size=self.e_shape)[np.newaxis].astype(
                self.dtype
            )
        else:
            edge_attr = None

        x = self.np_random.normal(size=self.x_shape)[np.newaxis].astype(self.dtype)

        return {
            "x": x,
            "edge_attr": edge_attr,
            "edge_index": np.array([[0], [0]]),
        }

    def contains(self, x):
        # placeholder
        return True

    def __repr__ (self):
        return "Graph({}, {})".format(self.x_shape, self.e_shape)

```

It does work in principle, however, there are a few bits that need to be adjusted.  
E.g. for batching RLlib assumes arrays with the same dimensionality, but in this case that’s not true, so instead of a normal Numpy array (size: batch x features) it creates an object array which breaks slicing downstream and Pytorch tensor conversion.

I haven’t finished the conversion because there are a lot of places that need to be changed, but to me it seems like there’s no fundamental issue in allowing variable input sizes, only syntactic.

Also, [Pytorch Geometric](https://pytorch-geometric.readthedocs.io/en/latest/) has a very convenient [batching](https://pytorch-geometric.readthedocs.io/en/latest/modules/data.html#torch_geometric.data.Batch) mechanism to address the same problem, while making it possible to handle the input tensors as a regular 2d tensor with size BN x F (batch size \* number of items x feature size). Integrating this approach would be a large endeavour, but it could enable a lot of applications that are currently limited by the current implementation.

So, is there any plan to enable more efficient and flexible dynamic spaces? I.e. pretty much the same functionality that `Repeated` has, but without the overhead of zero-padding.

---

<div class="post-metadata">

**Author:** ![gjoliver](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/gjoliver/32/1490_2.png) [@gjoliver](https://discuss.ray.io/u/gjoliver)\
**Post date:** [February 14, 2022, 6:38pm UTC](https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014/2 "2022-02-14T18:38:42Z")

</div>

this does sounds like it will break a lot of assumptions in RLlib. but there’s currently no plan to support additional dynamic spaces in RLlib.  
I wonder if you can reuse some of the utils (loss fn, etc) from RLlib, and build a custom solution.

---

<div class="post-metadata">

**Author:** ![vakker00](https://avatars.discourse-cdn.com/v4/letter/v/b5e925/32.png) [@vakker00](https://discuss.ray.io/u/vakker00)\
**Post date:** [February 14, 2022, 9:15pm UTC](https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014/3 "2022-02-14T21:15:29Z")

</div>

Yeah, I could work out a solution for a few of the assumptions, but honestly I don’t know how deep the issue goes. I would think that regular tensors are quite fundamental to the RLlib data pipelines unfortunately. On the other hand, the issue (at least for me) is on the model input side, so the loss function, action distributions, etc. are not affected. That would change though if there was a need for dynamic action spaces.

For the observations I think there’s a way to get around it, which involves something similar to [RepeatedValues](https://github.com/ray-project/ray/blob/master/rllib/models/repeated_values.py#L8) with a matching preprocessor (see [RepeatedValuesPreprocessor](https://github.com/ray-project/ray/blob/7f1bacc7dc9caf6d0ec042e39499bbf1d9a7d065/rllib/models/preprocessors.py#L300)), that would handle the batching and stacking. Maybe each observation could include a dummy scalar value for the `obs_flat` bit that is expected by some parts of the experience replay and alike, but I’m not sure if that would work.

Worst case scenario, the environment could put the observation directly into the object store and only pass the reference to the pipeline and the model would retrieve the objects based on their reference.

Something like:

```python
# Environment
def step(action):
    ....
    object_ref = ray.put(obs)
    return object_ref, reward, done, info

# Policy
def forward(input_dict):
    obs = [ray.get(ref) for ref in input_dict["obs"]]

    ....
    # to tensor, to GPU, etc.

```

It feels quite hacky though.

---

<div class="post-metadata">

**Author:** ![marco](https://avatars.discourse-cdn.com/v4/letter/m/a698b9/32.png) [@marco](https://discuss.ray.io/u/marco)\
**Post date:** [June 18, 2022, 1:08pm UTC](https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014/4 "2022-06-18T13:08:26Z")

</div>

Did you ever figure out how to make this work by any chance?

---

<div class="post-metadata">

**Author:** ![vakker00](https://avatars.discourse-cdn.com/v4/letter/v/b5e925/32.png) [@vakker00](https://discuss.ray.io/u/vakker00)\
**Post date:** [June 19, 2022, 10:50am UTC](https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014/5 "2022-06-19T10:50:25Z")

</div>

No, I’m still using `RepeatedValues`, it’s a quite inefficient, but works for now.

---

<div class="post-metadata">

**Author:** ![Mattias\_De\_Charleroy](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/mattias_de_charleroy/32/3251_2.png) [@Mattias\_De\_Charleroy](https://discuss.ray.io/u/Mattias_De_Charleroy)\
**Post date:** [September 27, 2022, 2:03pm UTC](https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014/6 "2022-09-27T14:03:01Z")

</div>

Note that gym spaces now also have a dynamic graph space implementation. [gym/graph.py at master · openai/gym (github.com)](https://github.com/openai/gym/blob/master/gym/spaces/graph.py)

---

<div class="post-metadata">

**Author:** ![Anthony\_Jung](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/anthony_jung/32/3537_2.png) [@Anthony\_Jung](https://discuss.ray.io/u/Anthony_Jung)\
**Post date:** [December 9, 2022, 9:13am UTC](https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014/7 "2022-12-09T09:13:24Z")

</div>

[vakker00](https://discuss.ray.io/u/vakker00) I was really pleased to see someone going under same situation.

Like you’ve come up with, I tried to bypass this way(in the process hacking),  
(see the code snippet)

but `ray.get` blocks all other remote operations.

until ray team supports flexible observation space “adaptor”  
we might have to rather write to a file than exploiting ObjectStore.

Not sure if I’m missing sth.

## Env Side

```auto
class ObjRefSpace(spaces.Box):
    def __init__ (self):
        super(). __init__ (shape=(28,), low=0, high=255, dtype=np.uint8)

        obs = nx.read_gexf('qsort_multidigraph.gexf')
        obs = parse_graph(obs)
        obj_ref = ray.put(obs)
        arr = RllibWrapper.encrypt(obj_ref)

        self.dummy = arr

    def sample(self):
        return self.dummy

    def contains(self, x):
        return True

class RllibWrapper(Wrapper):
    def __init__ (self, env: Env) -> None:
        if not ray.is_initialized():
            raise RuntimeError('Ray must be initialized before using RllibWrapper')

        super(). __init__ (env)
        # Repeated Values too inefficient
        # So we bypass by storing object into ObjectStore
        self.observation_space = ObjRefSpace()

        # self.flag = True

        # TODO: check if ray server is up

    @staticmethod
    def encrypt(obj_ref: ray.ObjectRef) -> np.ndarray:
        # obj_ref.binary() == b'\x00\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\x01\x00\x00\x00\x03'

        arr = np.array([i for i in obj_ref.binary()], dtype=np.uint8)

        return arr

    @staticmethod
    def decrypt(arr: np.ndarray) -> ObjectRef:
        # obj = ObjectRef(arr.tobytes())
        obj = ObjectRef(arr.tobytes())
            
        return obj 

    def _process_obs(self, obs):
        obs = parse_graph(obs)

        # e.g. ObjectRef(00ffffffffffffffffffffffffffffffffffffff0100000003000000)
        obj_ref = ray.put(obs)

        arr = RllibWrapper.encrypt(obj_ref)

        return arr

    def reset(self, **kwargs):
        obs = self.env.reset(**kwargs)

        return self._process_obs(obs)

    def step(self, action):
        obs, rew, done, info = self.env.step(action)
        obs = self._process_obs(obs)

        return obs, rew, done, info

```

## Policy Side

```auto
    @override(TorchModelV2)
    def forward(self, input_dict, state, seq_lens):
        batch = []
        for encrypted_arr in input_dict['obs']:
            arr = encrypted_arr.detach().cpu().numpy()
            obj_ref = RllibWrapper.decrypt(arr)
            arr = ray.get(obj_ref)
            print(arr)

```

---

<div class="post-metadata">

**Author:** ![Anthony\_Jung](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/anthony_jung/32/3537_2.png) [@Anthony\_Jung](https://discuss.ray.io/u/Anthony_Jung)\
**Post date:** [December 9, 2022, 9:29am UTC](https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014/8 "2022-12-09T09:29:34Z")

</div>

[vakker00](https://discuss.ray.io/u/vakker00) Pleased to see someone going under same situation

Like you’ve come up with, I tried to bypass this way.

(See the code snippet below)

but as I’m using remote rollout workers, `ray.get` will somehow block other remote operations and hang forever

Until ray team supports flexible or dummy observation, maybe we could rather write observation to a file.

Just let me know if I missed sth.

### Env Side

```auto
class ObjRefSpace(spaces.Box):
    def __init__ (self):
        super(). __init__ (shape=(28,), low=0, high=255, dtype=np.uint8)

        obs = {some dummy}
        obs = parse_graph(obs)
        obj_ref = ray.put(obs)
        arr = RllibWrapper.encrypt(obj_ref)

        self.dummy = arr

    def sample(self):
        return self.dummy

    def contains(self, x):
        return True

class RllibWrapper(Wrapper):
    def __init__ (self, env: Env) -> None:
        if not ray.is_initialized():
            raise RuntimeError('Ray must be initialized before using RllibWrapper')

        super(). __init__ (env)
        # Repeated Values too inefficient
        # So we bypass by storing object into ObjectStore
        self.observation_space = ObjRefSpace()

    @staticmethod
    def encrypt(obj_ref: ray.ObjectRef) -> np.ndarray:
        arr = np.array([i for i in obj_ref.binary()], dtype=np.uint8)

        return arr

    @staticmethod
    def decrypt(arr: np.ndarray) -> ObjectRef:
        # obj = ObjectRef(arr.tobytes())
        obj = ObjectRef(arr.tobytes())

        return obj

    def _process_obs(self, obs):
        obs = parse_graph(obs)
        obj_ref = ray.put(obs)

        arr = RllibWrapper.encrypt(obj_ref)

        return arr

    def reset(self, **kwargs):
        obs = self.env.reset(**kwargs)

        return self._process_obs(obs)

    def step(self, action):
        obs, rew, done, info = self.env.step(action)
        obs = self._process_obs(obs)

        return obs, rew, done, info

```

### Policy Side

```auto
    @override(TorchModelV2)
    def forward(self, input_dict, state, seq_lens):
        batch = []
        for encrypted_arr in input_dict['obs']:
            arr = encrypted_arr.detach().cpu().numpy()
            obj_ref = RllibWrapper.decrypt(arr)
            arr = ray.get(obj_ref, timeout=3)

```

---

<div class="post-metadata">

**Author:** ![vakker00](https://avatars.discourse-cdn.com/v4/letter/v/b5e925/32.png) [@vakker00](https://discuss.ray.io/u/vakker00)\
**Post date:** [December 9, 2022, 10:54am UTC](https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014/9 "2022-12-09T10:54:21Z")

</div>

Hm, interesting, not sure why that blocks. Maybe the object references mess up somehow when you encrypt/decrypt?  
Does this work with a simple script without any of the RLlib mechanisms?  
E.g.:

```python
obj = list(range(10))
obj_ref = encrypt(ray.put(obj))
obj_back = ray.get(decrypt(obj_ref))
assert obj == obj_back

```

---

<div class="post-metadata">

**Author:** ![Anthony\_Jung](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/anthony_jung/32/3537_2.png) [@Anthony\_Jung](https://discuss.ray.io/u/Anthony_Jung)\
**Post date:** [December 9, 2022, 11:42am UTC](https://discuss.ray.io/t/efficient-set-and-graph-space-for-rl/5014/10 "2022-12-09T11:42:30Z")

</div>

yes it works, above snippet is test-driven code!

Now I’m working on file write/read -ing strategy,  
it’s getting really hacky ☹

and even if we get to work with object store  
I expect file I/O would be faster since object store requires network I/O

correct me if I’m wrong!

@ray asking the team,  
Can calling `ray.get` inside remote policy cause trouble?

p.s. I’m just a high-level API using (for distributed computing)  
researcher, not really familiar with blocking stuffs!

Thanks!  
Anthony