<commit_before># coding=utf-8
# Copyright 2018 The TF-Agents Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Environment wrappers.

Wrappers in this module can be chained to change the overall behaviour of an
environment in common ways.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import collections
import gin
import numpy as np
import six
import tensorflow as tf

from tf_agents.environments import py_environment
from tf_agents.specs import array_spec
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step as ts
from tensorflow.python.util import nest  # pylint:disable=g-direct-tensorflow-import  # TF internal


class PyEnvironmentBaseWrapper(py_environment.PyEnvironment):
  """PyEnvironment wrapper forwards calls to the given environment."""

  def __init__(self, env):
    super(PyEnvironmentBaseWrapper, self).__init__()
    self._env = env

  def __getattr__(self, name):
    """Forward all other calls to the base environment."""
    return getattr(self._env, name)

  @property
  def batched(self):
    return getattr(self._env, 'batched', False)

  @property
  def batch_size(self):
    return getattr(self._env, 'batch_size', None)

  def _reset(self):
    return self._env.reset()

  def _step(self, action):
    return self._env.step(action)

  def observation_spec(self):
    return self._env.observation_spec()

  def action_spec(self):
    return self._env.action_spec()

  def render(self, mode='rgb_array'):
    return self._env.render(mode)

  def wrapped_env(self):
    return self._env


@gin.configurable
class TimeLimit(PyEnvironmentBaseWrapper):
  """End episodes after specified number of steps."""

  def __init__(self, env, duration):
    super(TimeLimit, self).__init__(env)
    self._duration = duration
    self._num_steps = None

  def _reset(self):
    self._num_steps = 0
    return self._env.reset()

  def _step(self, action):
    if self._num_steps is None:
      return self.reset()

    time_step = self._env.step(action)

    self._num_steps += 1
    if self._num_steps >= self._duration:
      time_step = time_step._replace(step_type=ts.StepType.LAST)

    if time_step.is_last():
      self._num_steps = None

    return time_step


@gin.configurable
class ActionRepeat(PyEnvironmentBaseWrapper):
  """Repeates actions over n-steps while acummulating the received reward."""

  def __init__(self, env, times):
    """Creates an action repeat wrapper.

    Args:
      env: Environment to wrap.<commit_msg>Expose TimeLimit's duration property publicly.<commit_after># coding=utf-8
# Copyright 2018 The TF-Agents Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Environment wrappers.

Wrappers in this module can be chained to change the overall behaviour of an
environment in common ways.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import collections
import gin
import numpy as np
import six
import tensorflow as tf

from tf_agents.environments import py_environment
from tf_agents.specs import array_spec
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step as ts
from tensorflow.python.util import nest  # pylint:disable=g-direct-tensorflow-import  # TF internal


class PyEnvironmentBaseWrapper(py_environment.PyEnvironment):
  """PyEnvironment wrapper forwards calls to the given environment."""

  def __init__(self, env):
    super(PyEnvironmentBaseWrapper, self).__init__()
    self._env = env

  def __getattr__(self, name):
    """Forward all other calls to the base environment."""
    return getattr(self._env, name)

  @property
  def batched(self):
    return getattr(self._env, 'batched', False)

  @property
  def batch_size(self):
    return getattr(self._env, 'batch_size', None)

  def _reset(self):
    return self._env.reset()

  def _step(self, action):
    return self._env.step(action)

  def observation_spec(self):
    return self._env.observation_spec()

  def action_spec(self):
    return self._env.action_spec()

  def render(self, mode='rgb_array'):
    return self._env.render(mode)

  def wrapped_env(self):
    return self._env


@gin.configurable
class TimeLimit(PyEnvironmentBaseWrapper):
  """End episodes after specified number of steps."""

  def __init__(self, env, duration):
    super(TimeLimit, self).__init__(env)
    self._duration = duration
    self._num_steps = None

  def _reset(self):
    self._num_steps = 0
    return self._env.reset()

  def _step(self, action):
    if self._num_steps is None:
      return self.reset()

    time_step = self._env.step(action)

    self._num_steps += 1
    if self._num_steps >= self._duration:
      time_step = time_step._replace(step_type=ts.StepType.LAST)

    if time_step.is_last():
      self._num_steps = None

    return time_step

  @property
  def duration(self):
    return self._duration


@gin.configurable
class ActionRepeat(PyEnvironmentBaseWrapper):
  """Repeates actions over n-steps while acummulating the received reward."""

  def __init__(self, env, times):
    """Creates an action repeat wrapper.

    Args:
      env: Environment to wrap.
{'fraction_non_alphanumeric': 0.08632286995515695, 'fraction_numerical': 0.003523382447149263, 'mean_word_length': 4.139917695473251, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3379542', 'n_tokens_mistral': 1840, 'n_tokens_neox': 1769, 'n_words': 618}