<commit_before># coding=utf-8
# Copyright 2020 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
#
#     https://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.

"""An API for representing constraints."""

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

import abc
import functools
from typing import Callable, Iterable, Optional, Text
import six

import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import
import tensorflow_probability as tfp

from tf_agents.bandits.policies import loss_utils
from tf_agents.bandits.specs import utils as bandit_spec_utils
from tf_agents.policies import utils as policy_utilities
from tf_agents.trajectories import time_step as ts
from tf_agents.typing import types
from tf_agents.utils import common
from tf_agents.utils import nest_utils


@six.add_metaclass(abc.ABCMeta)
class BaseConstraint(tf.Module):
  """Abstract base class for representing constraints.

  The constraint class provides feasibility computation functionality for
  computing the probability of actions being feasible.
  """

  def __init__(
      self,
      time_step_spec: types.TimeStep,
      action_spec: types.BoundedTensorSpec,
      name: Optional[Text] = None):
    """Initialization of the BaseConstraint class.

    Args:
      time_step_spec: A `TimeStep` spec of the expected time_steps.
      action_spec: A nest of `BoundedTensorSpec` representing the actions.
      name: Python str name of this constraint.
    """
    super(BaseConstraint, self).__init__(name=name)
    if not isinstance(time_step_spec, ts.TimeStep):
      raise ValueError(
          'The `time_step_spec` must be an instance of `TimeStep`, but is `{}`.'
          .format(type(time_step_spec)))

    self._time_step_spec = time_step_spec
    self._action_spec = action_spec

  # Subclasses must implement these methods.
  @abc.abstractmethod
  def __call__(self,
               observation: types.NestedTensor,
               actions: Optional[types.Tensor] = None) -> types.Tensor:
    """Returns the probability of input actions being feasible."""


class NeuralConstraint(BaseConstraint):
  """Class for representing a trainable constraint using a neural network.

  This constraint class uses a neural network to compute the action feasibility.
  In this case, the loss function needs to be exposed for training the neural
  network weights, typically done by the agent that uses this constraint.
  """

  def __init__(
      self,
      time_step_spec: types.TimeStep,
      action_spec: types.BoundedTensorSpec,
      constraint_network: Optional[types.Network],
      error_loss_fn: types.LossFn = tf.compat.v1.losses.mean_squared_error,
      name: Optional[Text] = 'NeuralConstraint'):
    """Creates a trainable constraint using a neural network.

    Args:
      time_step_spec: A `TimeStep` spec of the expected time_steps.
      action_spec: A nest of `BoundedTensorSpec` representing the actions.
      constraint_network: An instance of `tf_agents.network.Network` used to
        provide estimates of action feasibility. The input structure should be
        consistent with the `observation_spec`. If the constraint network is
        not available at construction time, it can be set later on using the
        constraint_network setter.
      error_loss_fn: A function for computing the loss used to train the
        constraint network. The default is `tf.losses.mean_squared_error`.
      name: Python str name of this agent. All variables in this module will
        fall under that name. Defaults to the class name.
    """
    super(NeuralConstraint, self).__init__(
        time_step_spec,
        action_spec,
        name)

    self._num_actions = policy_utilities.get_num_actions_from_tensor_spec(
        action_spec)
    if constraint_network is not None:
      with self.name_scope:
        constraint_network.create_variables()
    self._constraint_network = constraint_network
    self._error_loss_fn = error_loss_fn

  @property
  def constraint_network(self):
    return self._constraint_network

  @constraint_network.setter
  def constraint_network(self, constraint_network):
    if constraint_network is not None:
      with self.name_scope:
        constraint_network.create_variables()
    self._constraint_network = constraint_network

  def initialize(self):
    """Returns an op to initialize the constraint."""
    tf.compat.v1.variables_initializer(self.variables)

  def compute_loss(self,
                   observations: types.NestedTensor,
                   actions: types.NestedTensor,
                   rewards: types.Tensor,
                   weights: Optional[types.Float] = None,<commit_msg>Add support to override the observation spec in trainable constraints.<commit_after># coding=utf-8
# Copyright 2020 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
#
#     https://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.

"""An API for representing constraints."""

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

import abc
import functools
from typing import Callable, Iterable, Optional, Text
import six

import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import
import tensorflow_probability as tfp

from tf_agents.bandits.policies import loss_utils
from tf_agents.bandits.specs import utils as bandit_spec_utils
from tf_agents.policies import utils as policy_utilities
from tf_agents.trajectories import time_step as ts
from tf_agents.typing import types
from tf_agents.utils import common
from tf_agents.utils import nest_utils


@six.add_metaclass(abc.ABCMeta)
class BaseConstraint(tf.Module):
  """Abstract base class for representing constraints.

  The constraint class provides feasibility computation functionality for
  computing the probability of actions being feasible.
  """

  def __init__(
      self,
      time_step_spec: types.TimeStep,
      action_spec: types.BoundedTensorSpec,
      name: Optional[Text] = None):
    """Initialization of the BaseConstraint class.

    Args:
      time_step_spec: A `TimeStep` spec of the expected time_steps.
      action_spec: A nest of `BoundedTensorSpec` representing the actions.
      name: Python str name of this constraint.
    """
    super(BaseConstraint, self).__init__(name=name)
    if not isinstance(time_step_spec, ts.TimeStep):
      raise ValueError(
          'The `time_step_spec` must be an instance of `TimeStep`, but is `{}`.'
          .format(type(time_step_spec)))

    self._time_step_spec = time_step_spec
    self._action_spec = action_spec

  # Subclasses must implement these methods.
  @abc.abstractmethod
  def __call__(self,
               observation: types.NestedTensor,
               actions: Optional[types.Tensor] = None) -> types.Tensor:
    """Returns the probability of input actions being feasible."""


class NeuralConstraint(BaseConstraint):
  """Class for representing a trainable constraint using a neural network.

  This constraint class uses a neural network to compute the action feasibility.
  In this case, the loss function needs to be exposed for training the neural
  network weights, typically done by the agent that uses this constraint.
  """

  def __init__(
      self,
      time_step_spec: types.TimeStep,
      action_spec: types.BoundedTensorSpec,
      constraint_network: Optional[types.Network],
      error_loss_fn: types.LossFn = tf.compat.v1.losses.mean_squared_error,
      name: Optional[Text] = 'NeuralConstraint'):
    """Creates a trainable constraint using a neural network.

    Args:
      time_step_spec: A `TimeStep` spec of the expected time_steps.
      action_spec: A nest of `BoundedTensorSpec` representing the actions.
      constraint_network: An instance of `tf_agents.network.Network` used to
        provide estimates of action feasibility. The input structure should be
        consistent with the `observation_spec`. If the constraint network is
        not available at construction time, it can be set later on using the
        constraint_network setter.
      error_loss_fn: A function for computing the loss used to train the
        constraint network. The default is `tf.losses.mean_squared_error`.
      name: Python str name of this agent. All variables in this module will
        fall under that name. Defaults to the class name.
    """
    super(NeuralConstraint, self).__init__(
        time_step_spec,
        action_spec,
        name)

    self._num_actions = policy_utilities.get_num_actions_from_tensor_spec(
        action_spec)
    if constraint_network is not None:
      with self.name_scope:
        constraint_network.create_variables()
    self._constraint_network = constraint_network
    self._error_loss_fn = error_loss_fn

  @property
  def constraint_network(self):
    return self._constraint_network

  @constraint_network.setter
  def constraint_network(self, constraint_network):
    if constraint_network is not None:
      with self.name_scope:
        constraint_network.create_variables()
    self._constraint_network = constraint_network

  @property
  def observation_spec(self):
    return self._time_step_spec.observation

  @observation_spec.setter
  def observation_spec(self, observation_spec):
    self._time_step_spec = ts.TimeStep(
        step_type=self._time_step_spec.step_type,
        reward=self._time_step_spec.reward,
        discount=self._time_step_spec.discount,
        observation=observation_spec)

  def initialize(self):
    """Returns an op to initialize the constraint."""
    tf.compat.v1.variables_initializer(self.variables)

  def compute_loss(self,
                   observations: types.NestedTensor,
                   actions: types.NestedTensor,
                   rewards: types.Tensor,
                   weights: Optional[types.Float] = None,
{'fraction_non_alphanumeric': 0.06524153248195447, 'fraction_numerical': 0.0020359059781602813, 'mean_word_length': 3.6183760683760684, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 2, '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': '2967092', 'n_tokens_mistral': 2910, 'n_tokens_neox': 2842, 'n_words': 1015}