def METHOD_NAME(self, action):
    action_clipped = action
    if self.clamp_actions:
        action_clipped = torch.clamp(action, -1, 1)
    action_clipped = torch_to_jax(action_clipped)
    next_obs, reward, terminated, info = self.env.METHOD_NAME(action_clipped)
    next_obs = jax_to_torch(next_obs)
    reward = jax_to_torch(reward)
    terminated = jax_to_torch(terminated).to(torch.bool)
    truncated = jax_to_torch(info["truncation"]).to(torch.bool)
    if self.clamp_rew_obs:
        reward = torch.clamp(reward, -100, 100)
        next_obs = torch.clamp(next_obs, -100, 100)
    return next_obs, reward, terminated, truncated, info