Fitting Graph NNs into the Tianshou framework

[mjmartis]

Hi all,

I'm using Tianshou with a RL problem where the environment represents a network and the agent acts by modifying edges of the network graph.

The most effective models for graph problems are "GNN"s, many of which only accept examples in a particular form. Here's an example with Torch Geometric. Separate vectors are provided for the edge "coordinates", edge features and node features. Graph examples (even of different sizes) are batched together through a matrix diagonalisation trick.

The naive approach would be to output this structured data as the observation from my Env.step method, but this isn't possible (I see there is a Dict space that isn't supported by Tianshou). I could return it in the info dict for each step, however that doesn't seem like the intended use of info (e.g. it's missing when training on "next batch" observations in the DQN policy, but that might be an oversight).

What's the idiomatic way to pass (edge coordinate, edge feature, node feature) vectors of variable sizes to a GNN in the Tianshou framework?

Thanks

[mjmartis]

Here's a minimal example demonstrating my issue:

from venv import logger
import tianshou as ts
import gymnasium as gym
import torch


class TestNet(ts.utils.net.common.Net):
  ''' Prints args provided to the forward method. '''

  def __init__(self, s_shape, a_shape, h_sizes):
    super().__init__(state_shape=s_shape, action_shape=a_shape, hidden_sizes=h_sizes)

  def forward(self, obs, state, info):
    logger.warning(f'Model received: obs={obs} state={state} info={info}')
    return super().forward(obs, state, info)


class TestEnv(gym.Env):
  ''' Uses a dict state space. '''

  def __init__(self):
    super().__init__()
    self.action_space = gym.spaces.Discrete(2)
    self.observation_space = gym.spaces.Dict({
        # It appears the value of the 'obs' key is what is passed to the model.
        'obs': gym.spaces.Box(low=0, high=1, shape=(2,)),
        'extra_data': gym.spaces.Box(low=0, high=1, shape=(2,)),
    })

  def reset(self, seed=None, options=None):
    super().reset(seed=seed)
    return self.observation_space.sample(), {}

  def step(self, action):
    return self.observation_space.sample(), 0, True, False, {}


def test():
  env = TestEnv()

  net = TestNet(s_shape=(2, ), a_shape=2, h_sizes=[16, 16])
  optim = torch.optim.Adam(net.parameters(), lr=1e-2)
  policy = ts.policy.DQNPolicy(model=net, optim=optim, action_space=env.action_space)

  collector = ts.data.Collector(policy, env)
  collector.reset()
  collector.collect(n_step=1)
  # Prints:
  #   Model received: obs=[[0.15468411 0.8790275 ]] state=None info=Batch()

  test_collector = ts.data.Collector(policy, env)
  test_collector.reset()

  ts.trainer.OffpolicyTrainer(
      policy=policy,
      train_collector=collector,
      test_collector=test_collector,
      max_epoch=1,
      step_per_epoch=1,
      step_per_collect=1,
      batch_size=1,
      episode_per_test=1,
  ).run()
  # Prints:
  #   Model received: obs=[[0.5437955  0.99715394]] state=None info=[None]

test()

The first print shows how my model can't access structured observations and the second shows why I can't include my structured data in the info dict.