# Path: meltingpot/python/utils/substrates/wrappers/base.py
# class Wrapper(dmlab2d.Environment):
#   def __init__(self, env):
#   def reset(self, *args, **kwargs) -> ...:
#   def step(self, *args, **kwargs) -> ...:
#   def reward_spec(self, *args, **kwargs) -> ...:
#   def discount_spec(self, *args, **kwargs) -> ...:
#   def observation_spec(self, *args, **kwargs) -> ...:
#   def action_spec(self, *args, **kwargs) -> ...:
#   def close(self, *args, **kwargs) -> ...:
#   def observation(self, *args, **kwargs) -> ...:
#   def events(self, *args, **kwargs) -> ...:
#   def list_property(self, *args, **kwargs) -> ...:
#   def write_property(self, *args, **kwargs) -> ...:
#   def read_property(self, *args, **kwargs) -> ...:


import functools
import dm_env
import immutabledict
import numpy as np
from typing import Mapping, Sequence, TypeVar, Union
from meltingpot.python.utils.substrates.wrappers import base
    except ValueError:
      raise ValueError(f'Action {action_index} ({action}) does not match '
                       f'action_spec ({action_spec}).')


def _immutable_action(
    action: Mapping[str, Numeric],
    action_spec: Mapping[str, dm_env.specs.Array],
) -> Mapping[str, np.ndarray]:
  """Returns an immutable action."""
  new_action = {}
  for key, value in action.items():
    if isinstance(value, np.ndarray):
      value = np.copy(value)
    else:
      value = np.array(value, dtype=action_spec[key].dtype)
    value.flags.writeable = False
    new_action[key] = value
  return immutabledict.immutabledict(new_action)


def _immutable_action_table(
    action_table: Sequence[Mapping[str, Numeric]],
    action_spec: Mapping[str, dm_env.specs.Array],
) -> Sequence[Mapping[str, np.ndarray]]:
  """Returns an immutable action table."""
  return tuple(
      _immutable_action(action, action_spec) for action in action_table)



# Path: meltingpot/python/utils/substrates/wrappers/base.py
# class Wrapper(dmlab2d.Environment):
#   def __init__(self, env):
#   def reset(self, *args, **kwargs) -> ...:
#   def step(self, *args, **kwargs) -> ...:
#   def reward_spec(self, *args, **kwargs) -> ...:
#   def discount_spec(self, *args, **kwargs) -> ...:
#   def observation_spec(self, *args, **kwargs) -> ...:
#   def action_spec(self, *args, **kwargs) -> ...:
#   def close(self, *args, **kwargs) -> ...:
#   def observation(self, *args, **kwargs) -> ...:
#   def events(self, *args, **kwargs) -> ...:
#   def list_property(self, *args, **kwargs) -> ...:
#   def write_property(self, *args, **kwargs) -> ...:
#   def read_property(self, *args, **kwargs) -> ...:


# Path: meltingpot/python/utils/substrates/wrappers/discrete_action_wrapper.py
import functools
import dm_env
import immutabledict
import numpy as np
from typing import Mapping, Sequence, TypeVar, Union
from meltingpot.python.utils.substrates.wrappers import base

    except ValueError:
      raise ValueError(f'Action {action_index} ({action}) does not match '
                       f'action_spec ({action_spec}).')


def _immutable_action(
    action: Mapping[str, Numeric],
    action_spec: Mapping[str, dm_env.specs.Array],
) -> Mapping[str, np.ndarray]:
  """Returns an immutable action."""
  new_action = {}
  for key, value in action.items():
    if isinstance(value, np.ndarray):
      value = np.copy(value)
    else:
      value = np.array(value, dtype=action_spec[key].dtype)
    value.flags.writeable = False
    new_action[key] = value
  return immutabledict.immutabledict(new_action)


def _immutable_action_table(
    action_table: Sequence[Mapping[str, Numeric]],
    action_spec: Mapping[str, dm_env.specs.Array],
) -> Sequence[Mapping[str, np.ndarray]]:
  """Returns an immutable action table."""
  return tuple(
      _immutable_action(action, action_spec) for action in action_table)



class Wrapper(base.Wrapper):