hunk.code_window:                 "The following unknown entries in the runtime_env dictionary "
                   f"will be ignored: {unknown_fields}. If you intended to use "
                   "them as plugins, they must be nested in the `plugins` field.")
           # TODO(architkulkarni) This is to make it easy for the worker caching
           # code in C++ to check if the env is empty without deserializing and
           # parsing it.  We should use a less confusing approach here.
           if all(val is None for val in self.values()):
  
hunk.labels: keep keep keep keep replace replace replace keep
hunk.after_edit:         # NOTE(architkulkarni): This allows worker caching code in C++ to check
           # if a runtime env is empty without deserializing it.  This is a catch-
           # all; for validated inputs we won't set the key if the value is None.
  
hunk.file_path: python/ray/_private/runtime_env/validation.py
hunk.type: replace
file: from copy import deepcopy
  
  import numpy as np
  
  
  class RankedRewardsBuffer:
      def __init__(self, buffer_max_length, percentile):
          self.buffer_max_length = buffer_max_length
          self.percentile = percentile
          self.buffer = []
  
      def add_reward(self, reward):
          if len(self.buffer) < self.buffer_max_length:
              self.buffer.append(reward)
          else:
              self.buffer = self.buffer[1:] + [reward]
  
      def normalize(self, reward):
          reward_threshold = np.percentile(self.buffer, self.percentile)
          if reward < reward_threshold:
              return -1.0
          else:
              return 1.0
  
      def get_state(self):
          return np.array(self.buffer)
  
      def set_state(self, state):
          if state is not None:
              self.buffer = list(state)
  
  
  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 = self.env.reset()
                  done = False
                  while not done:
                      mask = obs["action_mask"]
                      probs = mask / mask.sum()
                      action = np.random.choice(
                          np.arange(mask.shape[0]), p=probs)
                      obs, reward, done, _ = self.env.step(action)
                  self.r2_buffer.add_reward(reward)
  
          def step(self, action):
              obs, reward, done, info = self.env.step(action)
              if done:
                  reward = self.r2_buffer.normalize(reward)
              return obs, reward, done, 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):
              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
  