forked from SeAIPalette/SeAIPalette
24 lines
795 B
Python
24 lines
795 B
Python
from Palette.algos.charge import Charge
|
|
|
|
|
|
class MultiCharge(object):
|
|
def __init__(self, n_ships: int, filling_time: int):
|
|
super().__init__()
|
|
self._n_ships = n_ships
|
|
self._filling_time = filling_time
|
|
self._charges = [Charge(self._filling_time)
|
|
for _ in range(self._n_ships)]
|
|
|
|
def __len__(self):
|
|
return self._n_ships
|
|
|
|
def __getitem__(self, idx):
|
|
return self._charges[idx]
|
|
|
|
def step(self, state, info, idx):
|
|
agent_info = {k: info[k]
|
|
for k in (info.keys()-{'fields', 'redecide_direction'})}
|
|
agent_info['fields'] = info['fields'][idx]
|
|
agent_info['redecide_direction'] = info['redecide_direction'][idx]
|
|
return self._charges[idx].step(state[idx], agent_info)
|