# 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) -> ...:


from typing import Iterator, Mapping, Sequence, TypeVar
from meltingpot.python.utils.substrates.wrappers import base
import dm_env
import numpy as np
# See the License for the specific language governing permissions and
# limitations under the License.
"""Wrapper that converts the DMLab2D specs into lists of action/observation."""




T = TypeVar("T")


def _player_observations(observations: Mapping[str, T], suffix: str,
                         num_players: int) -> Iterator[T]:
  """Yields observations for each player.

  Args:
    observations: dmlab2d observations source to check.
    suffix: suffix of player key to return.
    num_players: the number of players.
  """
  for player_index in range(num_players):
    try:
      value = observations[f"{player_index + 1}.{suffix}"]
    except KeyError:
      pass
    else:
      if isinstance(value, dm_env.specs.Array):
        value = value.replace(name=suffix)
      yield player_index, value



# 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/multiplayer_wrapper.py
from typing import Iterator, Mapping, Sequence, TypeVar
from meltingpot.python.utils.substrates.wrappers import base
import dm_env
import numpy as np

# See the License for the specific language governing permissions and
# limitations under the License.
"""Wrapper that converts the DMLab2D specs into lists of action/observation."""




T = TypeVar("T")


def _player_observations(observations: Mapping[str, T], suffix: str,
                         num_players: int) -> Iterator[T]:
  """Yields observations for each player.

  Args:
    observations: dmlab2d observations source to check.
    suffix: suffix of player key to return.
    num_players: the number of players.
  """
  for player_index in range(num_players):
    try:
      value = observations[f"{player_index + 1}.{suffix}"]
    except KeyError:
      pass
    else:
      if isinstance(value, dm_env.specs.Array):
        value = value.replace(name=suffix)
      yield player_index, value



class Wrapper(base.Wrapper):