Source code for rl_mind.envs

"""Extra environments.

`CartPoleContinuous-v1` is the classic cart-pole environment with a
*continuous* action in :math:`[-1, 1]` (the force applied to the cart), used
by the DDPG/TD3/SAC notebooks. Adapted from ``bbrl_gymnasium`` (MIT license,
Sorbonne Université).
"""

import gymnasium as gym
import numpy as np
from gymnasium import spaces
from gymnasium.envs.classic_control.cartpole import CartPoleEnv

__all__ = ["ContinuousCartPoleEnv"]




class ContinuousCartPoleEnv(CartPoleEnv):
    """Continuous-action version of the CartPole-v1 environment

    Same observation space, reward (+1 per step) and termination condition as
    `CartPole-v1`, but the action is the *signed fraction* of the maximum
    force applied to the cart, a `Box(-1, 1, shape=(1,))`. Registered as
    `CartPoleContinuous-v1` (500 steps at most) when `rl_mind.envs` is
    imported.
    """

    def __init__(self, **kwargs):
        """:param kwargs: Passed to `gymnasium`'s `CartPoleEnv` (e.g.
        `render_mode`)
        """
        super().__init__(**kwargs)
        self.action_space = spaces.Box(-1.0, 1.0, shape=(1,), dtype=np.float64)



    def step(self, action):
        """Apply a force to the cart and advance the simulation by one step

        :param action: The force, as a fraction of `force_mag` in $[-1, 1]$
            (clipped); a scalar or a one-element array
        :returns: The gymnasium 5-tuple `(obs, reward, terminated, truncated,
            info)`, with a reward of 1.0 at every step
        """
        force_scale = float(np.clip(np.asarray(action), -1.0, 1.0).item())
        assert self.state is not None, "Call reset before using step method."

        x, x_dot, theta, theta_dot = self.state
        force = self.force_mag * force_scale

        costheta = np.cos(theta)
        sintheta = np.sin(theta)

        # Dynamics: https://coneural.org/florian/papers/05_cart_pole.pdf
        temp = (
            force + self.polemass_length * theta_dot**2 * sintheta
        ) / self.total_mass
        thetaacc = (self.gravity * sintheta - costheta * temp) / (
            self.length * (4.0 / 3.0 - self.masspole * costheta**2 / self.total_mass)
        )
        xacc = temp - self.polemass_length * thetaacc * costheta / self.total_mass

        if self.kinematics_integrator == "euler":
            x = x + self.tau * x_dot
            x_dot = x_dot + self.tau * xacc
            theta = theta + self.tau * theta_dot
            theta_dot = theta_dot + self.tau * thetaacc
        else:  # semi-implicit euler
            x_dot = x_dot + self.tau * xacc
            x = x + self.tau * x_dot
            theta_dot = theta_dot + self.tau * thetaacc
            theta = theta + self.tau * theta_dot

        self.state = (x, x_dot, theta, theta_dot)

        terminated = bool(
            x < -self.x_threshold
            or x > self.x_threshold
            or theta < -self.theta_threshold_radians
            or theta > self.theta_threshold_radians
        )

        if self.render_mode == "human":
            self.render()
        return np.array(self.state, dtype=np.float32), 1.0, terminated, False, {}




# Importing this module registers the environments (an explicit
# `import rl_mind.envs` is the opt-in; the rl_mind package itself does not
# import it)
if "CartPoleContinuous-v1" not in gym.registry:
    gym.register(
        id="CartPoleContinuous-v1",
        entry_point=ContinuousCartPoleEnv,
        max_episode_steps=500,
    )