mirror of
https://github.com/illiumst/marl-factory-grid.git
synced 2025-05-23 07:16:44 +02:00
53 lines
1.8 KiB
Python
53 lines
1.8 KiB
Python
from typing import Union, List, Tuple
|
|
|
|
from marl_factory_grid.environment import constants as c
|
|
from marl_factory_grid.environment.groups.collection import Collection
|
|
from marl_factory_grid.modules.batteries.entitites import ChargePod, Battery
|
|
from marl_factory_grid.utils.results import Result
|
|
|
|
|
|
class Batteries(Collection):
|
|
_entity = Battery
|
|
|
|
@property
|
|
def var_has_position(self):
|
|
return False
|
|
|
|
@property
|
|
def var_can_be_bound(self):
|
|
return True
|
|
|
|
def __init__(self, size, initial_charge_level=1.0, *args, **kwargs):
|
|
"""
|
|
A collection of batteries that can spawn batteries.
|
|
|
|
:param size: The maximum allowed size of the collection. Ensures that the collection does not exceed this size.
|
|
:type size: int
|
|
|
|
:param initial_charge_level: The initial charge level of the battery.
|
|
:type initial_charge_level: float
|
|
"""
|
|
super(Batteries, self).__init__(size, *args, **kwargs)
|
|
self.initial_charge_level = initial_charge_level
|
|
|
|
def spawn(self, coords_or_quantity: Union[int, List[Tuple[(int, int)]]], *entity_args, **entity_kwargs):
|
|
batteries = [self._entity(self.initial_charge_level, agent) for _, agent in enumerate(entity_args[0])]
|
|
self.add_items(batteries)
|
|
|
|
def trigger_spawn(self, state, *entity_args, coords_or_quantity=None, **entity_kwargs):
|
|
self.spawn(0, state[c.AGENT])
|
|
return Result(identifier=f'{self.name}_spawn', validity=c.VALID, value=len(self))
|
|
|
|
|
|
class ChargePods(Collection):
|
|
_entity = ChargePod
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
"""
|
|
A collection of charge pods in the environment.
|
|
"""
|
|
super(ChargePods, self).__init__(*args, **kwargs)
|
|
|
|
def __repr__(self):
|
|
return super(ChargePods, self).__repr__()
|