Replay memory for Graph Data! TypeError: expected Tensor as element 0 in argument 0, but got Data

Famirreza: <p>I have this issue when I am trying to use <a href="https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html#training-loop" rel="noopener nofollow ugc">optimize_model()</a> in my DQN training process for <a href="https://pytorch-geometric.readthedocs.io/en/latest/modules/data.html#torch_geometric.data.Data" rel="noopener nofollow ugc">torch graph data type</a>.</p>
<p>Could you please explain how should I use this function for my torch graph data?</p>
<pre><code class="lang-auto">def optimize_model():
    if len(memory) &lt; BATCH_SIZE:
        return
    transitions = memory.sample(BATCH_SIZE)
    # Transpose the batch (see https://stackoverflow.com/a/19343/3343043 for
    # detailed explanation). This converts batch-array of Transitions
    # to Transition of batch-arrays.
    batch = Transition(*zip(*transitions))

    #print("Batch is here:", batch)
    # Compute a mask of non-final states and concatenate the batch elements
    # (a final state would've been the one after which simulation ended)


    non_final_mask = torch.tensor(tuple(map(lambda s: s is not None, batch.next_state)), device=device, dtype=torch.bool)    
    non_final_next_states =  torch.cat(([s for s in batch.next_state if s is not None]))
    
    state_batch =  torch.cat(batch.state)
    action_batch =  torch.cat(batch.action)
    reward_batch =  torch.cat(batch.reward)


    # Compute Q(s_t, a) - the model computes Q(s_t), then we select the
    # columns of actions taken. These are the actions which would've been taken
    # for each batch state according to policy_net
    state_action_values = policy_net(state_batch.x, state_batch.edge_index, state_batch.edge_type, state_batch.edge_attr).gather(1, action_batch)

</code></pre>

Famirreza: <p>I also tried to the resolve issue using splitting graph data to different components (please check below) but still I have AssertionError!</p>
<pre><code class="lang-auto">def optimize_model():
    if len(memory) &lt; BATCH_SIZE:
        return
    transitions = memory.sample(BATCH_SIZE)
    # Transpose the batch (see https://stackoverflow.com/a/19343/3343043 for
    # detailed explanation). This converts batch-array of Transitions
    # to Transition of batch-arrays.
    batch = Transition(*zip(*transitions))

    #print("Batch is here:", batch)
    # Compute a mask of non-final states and concatenate the batch elements
    # (a final state would've been the one after which simulation ended)


    non_final_mask = torch.tensor(tuple(map(lambda s: s is not None, batch.next_state)), device=device, dtype=torch.bool)    

    non_final_next_states_x_s = torch.cat(([s.x for s in batch.next_state if s is not None]))
    non_final_next_states_edge_index_s = torch.cat(([s.edge_index for s in batch.next_state if s is not None])) 
    non_final_next_states_edge_type_s = torch.cat(([s.edge_type for s in batch.next_state if s is not None])) 
    non_final_next_states_edge_attr_s = torch.cat(([s.edge_attr for s in batch.next_state if s is not None])) 

    x_s = torch.cat(([s.x for s in batch.state]))
    edge_index_s = torch.cat(([s.edge_index for s in batch.state])) 
    edge_type_s = torch.cat(([s.edge_type for s in batch.state])) 
    edge_attr_s = torch.cat(([s.edge_attr for s in batch.state]))

    action_batch =  torch.tensor(batch.action)
    reward_batch =  torch.tensor(batch.reward)
    

    # Compute Q(s_t, a) - the model computes Q(s_t), then we select the
    # columns of actions taken. These are the actions which would've been taken
    # for each batch state according to policy_net
    state_action_values = policy_net(x_s, edge_index_s, edge_type_s, edge_attr_s).gather(1, action_batch)

</code></pre>

vmoens: <p>Thanks for reporting this. Can I ask you a bit more context, e.g. a minimally reproducible example, the exact error message and what you are trying to achieve?<br>
From a high level perspective you are trying to use a replay buffer that consumes tensors with a data structure that is not a tensor.<br>
If you’d like <a href="https://github.com/facebookresearch/rl" rel="noopener nofollow ugc">torchrl</a> provides <a href="https://github.com/facebookresearch/rl/tree/main/torchrl/data/replay_buffers" rel="noopener nofollow ugc">replay buffers</a> that do not assume any data type. Have a look at the <a href="https://github.com/facebookresearch/rl/blob/main/tutorials/demo.ipynb" rel="noopener nofollow ugc">demo</a> for some pointers.</p>
<p>Happy to help further once we get some more info</p>