# Path: deluca/core.py
# class Env(Obj):
# 
#   @abstractmethod
#   def init(self):
#     """Return an initialized state"""
# 
#   @abstractmethod
#   def __call__(self, state, action, *args, **kwargs):
#     """Return an updated state"""
# 
# Path: deluca/core.py
# def field(default=None, jaxed=True, **kwargs):
#   if "default_factory" not in kwargs:
#     kwargs["default"] = default
#   kwargs["pytree_node"] = jaxed
#   return flax.struct.field(**kwargs)


from deluca.core import Env
from deluca.core import field
import jax
import jax.numpy as jnp
    "Jonathan P. How",
]
__license__ = "BSD 3-Clause"
__author__ = "Christoph Dann <cdann@cdann.de>"

# SOURCE:
# https://github.com/rlpy/rlpy/blob/master/rlpy/Domains/Acrobot.py


class Acrobot(Env):
  """Acrobot."""
  key: jnp.ndarray = field(jaxed=False)
  dt: float = field(0.2, jaxed=False)

  LINK_LENGTH_1: float = field(1.0, jaxed=False)
  LINK_LENGTH_2: float = field(1.0, jaxed=False)
  LINK_MASS_1: float = field(1.0, jaxed=False)
  LINK_MASS_2: float = field(1.0, jaxed=False)
  LINK_COM_POS_1: float = field(0.0, jaxed=False)
  LINK_COM_POS_2: float = field(0.0, jaxed=False)
  LINK_MOI: float = field(1.0)

  MAX_VEL_1: float = field(4 * jnp.pi, jaxed=False)
  MAX_VEL_2: float = field(9 * jnp.pi, jaxed=False)

  AVAIL_TORQUE: jnp.ndarray = field(jaxed=False)

  torque_noise_max: float = field(0.0, jaxed=False)

  high: jnp.ndarray = field(jaxed=False)

# Path: deluca/core.py
# class Env(Obj):
# 
#   @abstractmethod
#   def init(self):
#     """Return an initialized state"""
# 
#   @abstractmethod
#   def __call__(self, state, action, *args, **kwargs):
#     """Return an updated state"""
# 
# Path: deluca/core.py
# def field(default=None, jaxed=True, **kwargs):
#   if "default_factory" not in kwargs:
#     kwargs["default"] = default
#   kwargs["pytree_node"] = jaxed
#   return flax.struct.field(**kwargs)


# Path: deluca/envs/classic/_acrobot.py
from deluca.core import Env
from deluca.core import field
import jax
import jax.numpy as jnp

    "Jonathan P. How",
]
__license__ = "BSD 3-Clause"
__author__ = "Christoph Dann <cdann@cdann.de>"

# SOURCE:
# https://github.com/rlpy/rlpy/blob/master/rlpy/Domains/Acrobot.py


class Acrobot(Env):
  """Acrobot."""
  key: jnp.ndarray = field(jaxed=False)
  dt: float = field(0.2, jaxed=False)

  LINK_LENGTH_1: float = field(1.0, jaxed=False)
  LINK_LENGTH_2: float = field(1.0, jaxed=False)
  LINK_MASS_1: float = field(1.0, jaxed=False)
  LINK_MASS_2: float = field(1.0, jaxed=False)
  LINK_COM_POS_1: float = field(0.0, jaxed=False)
  LINK_COM_POS_2: float = field(0.0, jaxed=False)
  LINK_MOI: float = field(1.0)

  MAX_VEL_1: float = field(4 * jnp.pi, jaxed=False)
  MAX_VEL_2: float = field(9 * jnp.pi, jaxed=False)

  AVAIL_TORQUE: jnp.ndarray = field(jaxed=False)

  torque_noise_max: float = field(0.0, jaxed=False)

  high: jnp.ndarray = field(jaxed=False)

  low: jnp.ndarray = field(jaxed=False)