import torch
import torch.nn as nn
import torch.nn.functional as F

class RNNPolicyNetwork(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(RNNPolicyNetwork, self).__init__()
        # Define an LSTM layer
        self.rnn = nn.LSTM(input_dim, 128, batch_first=True)
        self.fc = nn.Linear(128, output_dim)

    def forward(self, x, hx):
        # Pass input through the LSTM layer
        x, hx = self.rnn(x, hx)
        # Use the output from the last time step
        x = F.softmax(self.fc(x[:, -1, :]), dim=-1)
        return x, hx