code: def get_r2_env_wrapper(env_creator, r2_config):
      class RankedRewardsEnvWrapper:
          def __init__(self, env_config):
              self.env = env_creator(env_config)
              self.action_space = self.env.action_space
              self.observation_space = self.env.observation_space
              max_buffer_length = r2_config["buffer_max_length"]
              percentile = r2_config["percentile"]
              self.r2_buffer = RankedRewardsBuffer(
                  max_buffer_length, percentile
              )
              if r2_config["initialize_buffer"]:
                  self._initialize_buffer(r2_config["num_init_rewards"])
  
          def _initialize_buffer(self, num_init_rewards=100):
              # initialize buffer with random policy
              for _ in range(num_init_rewards):
                  obs, info = self.env.reset()
                  terminated = truncated = False
                  while not terminated and not truncated:
                      mask = obs["action_mask"]
                      probs = mask / mask.sum()
                      action = np.random.choice(
                          np.arange(mask.shape[0]), p=probs
                      )
                      (
                          obs,
                          reward,
                          terminated,
                          truncated,
                          _,
                      ) = self.env.step(action)
                  self.r2_buffer.add_reward(reward)
  
          def step(self, action):
              obs, reward, terminated, truncated, info = self.env.step(
                  action
              )
              if terminated or truncated:
                  reward = self.r2_buffer.normalize(reward)
              return obs, reward, terminated, truncated, info
  
          def get_state(self):
              state = {
                  "env_state": self.env.get_state(),
                  "buffer_state": self.r2_buffer.get_state(),
              }
              return deepcopy(state)
  
          def reset(self, *, seed=None, options=None):
              return self.env.reset()
  
          def set_state(self, state):
              obs = self.env.set_state(state["env_state"])
              self.r2_buffer.set_state(state["buffer_state"])
              return obs
  
      return RankedRewardsEnvWrapper
  
file_path: ray/rllib_contrib/alpha_zero/src/rllib_alpha_zero/alpha_zero/ranked_rewards.py
imports: from copy import deepcopy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: get_r2_env_wrapper(env_creator, r2_config) and returns: RankedRewardsEnvWrapper, obs,  reward,  terminated,  truncated,  info, deepcopy, self, obs