hunk.code_window:     for a_name, action_space in ACTION_SPACES_TO_TEST.items():
           for o_name, obs_space in OBSERVATION_SPACES_TO_TEST.items():
               print("=== Testing {} A={} S={} ===".format(
                   alg, action_space, obs_space))
               stub_env = make_stub_env(action_space, obs_space, check_bounds)
               register_env("stub_env", lambda c: stub_env())
               stat = "ok"
               a = None
               try:
  
hunk.labels: keep keep replace replace keep keep keep keep keep
hunk.after_edit:             print("=== Testing {} (torch={}) A={} S={} ===".format(
                   alg, torch, action_space, obs_space))
  
hunk.file_path: rllib/tests/test_supported_spaces.py
hunk.type: replace
file: import numpy as np
  
  from ray.rllib.models.action_dist import ActionDistribution
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils import try_import_torch
  
  torch, nn = try_import_torch()
  
  
  class TorchDistributionWrapper(ActionDistribution):
      """Wrapper class for torch.distributions."""
  
      def __init_(self, inputs):
          super().__init__(inputs)
          # Store the last sample here.
          self.last_sample = None
  
      @override(ActionDistribution)
      def logp(self, actions):
          return self.dist.log_prob(actions)
  
      @override(ActionDistribution)
      def entropy(self):
          return self.dist.entropy()
  
      @override(ActionDistribution)
      def kl(self, other):
          return torch.distributions.kl.kl_divergence(self.dist, other.dist)
  
      @override(ActionDistribution)
      def sample(self):
          self.last_sample = self.dist.sample()
          return self.last_sample
  
      @override(ActionDistribution)
      def sampled_action_logp(self):
          assert self.last_sample is not None
          return self.logp(self.last_sample)
  
  
  class TorchCategorical(TorchDistributionWrapper):
      """Wrapper class for PyTorch Categorical distribution."""
  
      @override(ActionDistribution)
      def __init__(self, inputs, model):
          super().__init__(inputs, model)
          self.dist = torch.distributions.categorical.Categorical(logits=inputs)
  
      @override(ActionDistribution)
      def deterministic_sample(self):
          return self.dist.probs.argmax(dim=1)
  
      @staticmethod
      @override(ActionDistribution)
      def required_model_output_shape(action_space, model_config):
          return action_space.n
  
  
  class TorchDiagGaussian(TorchDistributionWrapper):
      """Wrapper class for PyTorch Normal distribution."""
  
      @override(ActionDistribution)
      def __init__(self, inputs, model):
          super().__init__(inputs, model)
          mean, log_std = torch.chunk(inputs, 2, dim=1)
          self.dist = torch.distributions.normal.Normal(mean, torch.exp(log_std))
  
      @override(ActionDistribution)
      def deterministic_sample(self):
          return self.dist.mean
  
      @override(TorchDistributionWrapper)
      def logp(self, actions):
          return TorchDistributionWrapper.logp(self, actions).sum(-1)
  
      @staticmethod
      @override(ActionDistribution)
      def required_model_output_shape(action_space, model_config):
          return np.prod(action_space.shape) * 2
  