Implementing a custom RNN using the TorchModelV2

Kunzro · December 16, 2022, 3:08pm

What exactly is the expected return (shape and to what do the dimensions correspond to, also what if the batch_size is only 1) from the forward function of the TorchModelV2 if it is used to implement a RNN. I have the same question regarding the value_function(). First and foremost: Do I have to keep the “Sequence” dimension for the return values in either of the forward() or value_function() function?

This is possible according to the documentation:

Note that the inputs arg entering forward_rnn is already a time-ranked single tensor (not an input_dict!) with shape (B x T x ...). If you further want to customize and need more direct access to the complete (non time-ranked) input_dict, you can also override your Model’s forward method directly (as you would do with a non-RNN ModelV2). In that case, though, you are responsible for changing your inputs and add the time rank to the incoming data (usually you just have to reshape).

I want to use the “normal” TorchModelV2 since I want to have more control over the inputs, since my state consists of a Dict and I want to handle individual parts differently, while the TorchRNN model just flattens everything and concatenates it.

How severe does this issue affect your experience of using Ray?

High: It blocks me to complete my task.

mannyv · December 16, 2022, 4:17pm

Hi @Kunzro,

Here is an example they might help you get started.

github.com

ray-project/ray/blob/master/rllib/examples/models/rnn_model.py

import numpy as np

from ray.rllib.models.modelv2 import ModelV2
from ray.rllib.models.preprocessors import get_preprocessor
from ray.rllib.models.tf.recurrent_net import RecurrentNetwork
from ray.rllib.models.torch.recurrent_net import RecurrentNetwork as TorchRNN
from ray.rllib.utils.annotations import override
from ray.rllib.utils.framework import try_import_tf, try_import_torch

tf1, tf, tfv = try_import_tf()
torch, nn = try_import_torch()


class RNNModel(RecurrentNetwork):
    """Example of using the Keras functional API to define a RNN model."""

    def __init__(
        self,
        obs_space,
        action_space,

This file has been truncated. show original

github.com

ray-project/ray/blob/master/rllib/examples/custom_rnn_model.py

"""Example of using a custom RNN keras model."""

import argparse
import os

import ray
from ray import air, tune
from ray.tune.registry import register_env
from ray.rllib.examples.env.repeat_after_me_env import RepeatAfterMeEnv
from ray.rllib.examples.env.repeat_initial_obs_env import RepeatInitialObsEnv
from ray.rllib.examples.models.rnn_model import RNNModel, TorchRNNModel
from ray.rllib.models import ModelCatalog
from ray.rllib.utils.test_utils import check_learning_achieved
from ray.tune.registry import get_trainable_cls

parser = argparse.ArgumentParser()
parser.add_argument(
    "--run", type=str, default="PPO", help="The RLlib-registered algorithm to use."
)
parser.add_argument("--env", type=str, default="RepeatAfterMeEnv")

This file has been truncated. show original

Topic		Replies	Views
Custom LSTM Model, how to define the SEQ_LEN RLlib	5	2352	June 10, 2024
State shapes incorrect using custom model (TorchModelV2) (PPO) RLlib	2	421	July 15, 2021
Input_dict's data structure RLlib	2	107	November 25, 2024
[RLlib] Shape Error for custom PyTorch model RLlib	2	685	March 12, 2021
Custom RNN Model with Examples - why do they fail? RLlib	11	2308	May 5, 2022

Implementing a custom RNN using the TorchModelV2

Related topics