import flax
from ml_collections import ConfigDict

from jax_rl.agents.common import Model, Params, PRNGKey


@flax.struct.dataclass
class SAC:
    actor: Model
    critic: Model
    target_critic: Model
    temp: Model
    config: ConfigDict = flax.struct.field(pytree_node=False)
    rng: PRNGKey