[Feature Request] CUDNN version of RSSM

Motivation

In model based models such as dreamer and planet the main computation bottleneck is the unroll of the RSSM prior. RSSM prior is based on a GRUCell. However, since at each step we need to compute some additional operations, we cannot use nn.GRU to have an optimized version. Having this could considerably boost the performance of model based model that use RSSM.

Solution

Currently, we have a single operation of RSSMPrior defined as:

class RSSMPrior(nn.Module):
    def __init__(self, hidden_dim=200, rnn_hidden_dim=200, state_dim=20):
        super().__init__()
        ### Prior
        self.rnn = nn.GRUCell(hidden_dim, rnn_hidden_dim)
        self.action_state_projector = nn.Sequential(nn.LazyLinear(hidden_dim), nn.ELU())
        self.rnn_to_prior_projector = NormalParamWrapper(
            nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim),
                nn.ELU(),
                nn.Linear(hidden_dim, 2 * state_dim),
            ),
            scale_lb=0.1,
        )

        self.state_dim = state_dim
        self.rnn_hidden_dim = rnn_hidden_dim

    def forward(self, state, rnn_hidden, action):
        *batch_size, _ = action.shape
        if state is None:
            state = torch.zeros(*batch_size, self.state_dim, device=action.device)
        if rnn_hidden is None:
            rnn_hidden = torch.zeros(
                *batch_size, self.rnn_hidden_dim, device=action.device
            )

        action_state = self.action_state_projector(torch.cat([state, action], dim=-1))
        rnn_hidden = self.rnn(action_state, rnn_hidden)
        belief = rnn_hidden
        prior_mean, prior_std = self.rnn_to_prior_projector(belief)
        prior_state = prior_mean + torch.randn_like(prior_std) * prior_std
        return prior_mean, prior_std, prior_state, belief

And what we unroll it in two ways:

1. Given a actions of shape [BxTxD], (init_prior_state, init_belief of shape [BxD1] and [BxD2] resp.) we compute the T corresponding states

    def unroll_rssm(init_prior_state, init_belief , actions):
        rssm_prior = RSSMPrior(hidden_dim=200, rnn_hidden_dim=200, state_dim=20)
        prior_means = []
        prior_stds = []
        prior_states = []
        beliefs = []
        prior_state, belief = init_prior_state, init_belief`
        for i in range(actions.shape[1]):
            prior_mean, prior_std, prior_state, belief = rssm_prior(
                prior_state, belief, action[:, i]
            )
            prior_means.append(prior_mean)
            prior_stds.append(prior_std)
            prior_states.append(prior_state)
            beliefs.append(belief)
        prior_means = torch.stack(prior_means, dim=1)
        prior_stds = torch.stack(prior_stds, dim=1)
        prior_states = torch.stack(prior_states, dim=1)
        beliefs = torch.stack(beliefs, dim=1)
        return prior_means, prior_stds, prior_states, beliefs

2. Given an actor model, we unroll the states for H steps computing the actions on the fly.

    def unroll_rssm(init_prior_state, init_belief , actor, horizon=H):
        rssm_prior = RSSMPrior(hidden_dim=200, rnn_hidden_dim=200, state_dim=20)
        prior_means = []
        prior_stds = []
        prior_states = []
        beliefs = []
        prior_state, belief = init_prior_state, init_belief
        for i in range(horizon):
            action = actor(prior_state, belief)
            prior_mean, prior_std, prior_state, belief = rssm_prior(
                prior_state, belief, action
            )
            prior_means.append(prior_mean)
            prior_stds.append(prior_std)
            prior_states.append(prior_state)
            beliefs.append(belief)
        prior_means = torch.stack(prior_means, dim=1)
        prior_stds = torch.stack(prior_stds, dim=1)
        prior_states = torch.stack(prior_states, dim=1)
        beliefs = torch.stack(beliefs, dim=1)
        return prior_means, prior_stds, prior_states, beliefs

The idea would be to have an optimized CUDNN version of both operation which would speed-up this operation.

Checklist

• I have checked that there is no similar issue in the repo (required)