from open_source.rlpyt.rlpyt.agents.dqn.r2d1_agent import R2d1Agent, R2d1AlternatingAgent
from open_source.rlpyt.rlpyt.models.dqn.atari_r2d1_model import AtariR2d1Model
from open_source.rlpyt.rlpyt.agents.dqn.atari.mixin import AtariMixin


class AtariR2d1Agent(AtariMixin, R2d1Agent):

    def __init__(self, ModelCls=AtariR2d1Model, **kwargs):
        super().__init__(ModelCls=ModelCls, **kwargs)


class AtariR2d1AlternatingAgent(AtariMixin, R2d1AlternatingAgent):

    def __init__(self, ModelCls=AtariR2d1Model, **kwargs):
        super().__init__(ModelCls=ModelCls, **kwargs)
    
