Project: 055_Prioritized_Generative_Replay | Function: compute_reward (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements intrinsic curiosity module (ICM) to compute relevance function F for prioritized generative replay. Uses forward dynamics model prediction error as curiosity-driven priority signal to guide conditional generation towards more learning-relevant transitions.

**KEY TERMS:** {
  "real_next_state_feature": "h(s′)",
  "pred_next_state_feature": "g(h(s), a)",
  "icm_reward": "F(s, a, s′, r)",
  "mse_loss": "1/2 ||·||^2"
}

**CONTEXT CODE:**             nn.Linear(self.output_size + self.hidden_size, self.hidden_size),
            Swish(),
            nn.Linear(self.hidden_size, self.hidden_size),
            Swish(),
            nn.Linear(self.hidden_size, self.hidden_size),
            Swish(),
            nn.Linear(self.hidden_size, self.hidden_size),
            Swish(),
            nn.Linear(self.hidden_size, self.hidden_size)
        )
        self.forward_net_2 = nn.Sequential(
            nn.Linear(self.output_size + self.hidden_size, self.hidden_size),
            Swish(),
            nn.Linear(self.hidden_size, self.hidden_size),
            Swish(),
            nn.Linear(self.hidden_size, self.hidden_size),
            Swish(),
            nn.Linear(self.hidden_size, self.hidden_size),
            Swish(),
            nn.Linear(self.hidden_size, self.hidden_size)
        )

    @property
    def device(self):
        return next(self.parameters()).device

    def forward(self, state, next_state, action):
        encode_state = self.feature(state)
        encode_next_state = self.feature(next_state)
        # get pred action
        pred_action = torch.cat((encode_state, encode_next_state), 1)
        pred_action = self.inverse_net(pred_action)
        # ---------------------

        # get pred next state
        pred_next_state_feature_orig = torch.cat((encode_state, action), 1)
        pred_next_state_feature_orig = self.forward_net_1(pred_next_state_feature_orig)

        # residual
        for i in range(self.resnet_time):
            pred_next_state_feature = self.residual[i * 2](torch.cat((pred_next_state_feature_orig, action), 1))
            pred_next_state_feature_orig = self.residual[i * 2 + 1](
                torch.cat((pred_next_state_feature, action), 1)) + pred_next_state_feature_orig

        pred_next_state_feature = self.forward_net_2(torch.cat((pred_next_state_feature_orig, action), 1))

        real_next_state_feature = encode_next_state

        return real_next_state_feature, pred_next_state_feature, pred_action

    def compute_reward(self, state, next_state, action, reward, done):
        raise NotImplementedError("Function implementation has been masked for testing")
    
    def compute_reward_torch(self, state, next_state, action):
        real_next_state_feature, pred_next_state_feature, _ = self.forward(state, next_state, action)
        icm_reward = F.mse_loss(real_next_state_feature, pred_next_state_feature, reduction='none').mean(1, keepdim=True)

        return icm_reward

    def forward_loss(self, state, next_state, action):
        real_next_state_feature, pred_next_state_feature, pred_action = self.forward(state, next_state, action)
        loss = F.mse_loss(real_next_state_feature.detach(), pred_next_state_feature) + F.mse_loss(action, pred_action)
        return loss



Complete the function implementation replacing "raise NotImplementedError".