Skip to content

Algorithm API

Dynamic programming

gym_classics2.algorithms.dynamic_programming

This file implements dynamic programming algorithms for solving Markov Decision Processes (MDPs) in gym-classics environments with model access. The algorithms include value iteration and policy iteration, which are fundamental methods in reinforcement learning for computing optimal policies and value functions.

backup

backup(env, discount, V, state, action)

Computes the Bellman backup for a given state and action.

Parameters:

Name Type Description Default
env

A gym-classics environment with model access.

required
discount

The discount factor.

required
V

The current value function.

required
state

The current state.

required
action

The action to evaluate.

required

Returns: The computed Q-value for the given state and action.

Source code in gym_classics2/algorithms/dynamic_programming.py
def backup(env, discount, V, state, action):
    """Computes the Bellman backup for a given state and action.

    Args:
        env: A gym-classics environment with model access.
        discount: The discount factor.
        V: The current value function.
        state: The current state.
        action: The action to evaluate.
    Returns:
        The computed Q-value for the given state and action.
    """

    V = np.array(V)

    next_states, rewards, terminals, probs = env.model(state, action)
    bootstraps = (1.0 - terminals) * V[next_states]
    return np.sum(probs * (rewards + discount * bootstraps))

value_iteration

value_iteration(env, discount, precision=0.001, history=False, verbose=False)

Performs value iteration for the given environment.

Parameters:

Name Type Description Default
env

A gym-classics environment with model access.

required
discount

The discount factor (0 <= discount <= 1).

required
precision

The precision for convergence (default: 1e-3).

0.001
history

If True, returns a list of intermediate value functions.

False
verbose

If True, prints progress information.

False

Returns:

Type Description

The optimal value function V. If history is True, returns a list of intermediate value functions.

Source code in gym_classics2/algorithms/dynamic_programming.py
def value_iteration(env, discount, precision=1e-3, history = False, verbose = False):
    """Performs value iteration for the given environment.

    Args:
        env: A gym-classics environment with model access.
        discount: The discount factor (0 <= discount <= 1).
        precision: The precision for convergence (default: 1e-3).
        history: If True, returns a list of intermediate value functions.
        verbose: If True, prints progress information.

    Returns:
        The optimal value function V. If history is True, returns a list of intermediate value functions.
    """

    assert isinstance(env, GymClassicsBaseEnv), "Value iteration requires a gym-classics environment with model access to the environment." 
    assert 0.0 <= discount <= 1.0
    assert precision > 0.0

    V = np.zeros(len(env.states()), dtype=np.float64)  
    if history:
        V_list = []
        V_list.append(V.copy())

    sweeps = 0
    progress = tqdm(total=None, desc="Value Iteration", disable=verbose)
    while True:
        progress.update()
        if verbose:
            print('.', end = '')
            sweeps += 1

        V_old = V.copy()

        for s in env.states():
            Q_values = [backup(env, discount, V, s, a) for a in range(len(env.actions()))]
            V[s] = np.max(Q_values)

        if history:
            V_list.append(V.copy())

        if np.abs(V - V_old).max() <= precision:
            break

    if verbose:
        print(f'\nConverged after {sweeps} sweeps.')

    if history:
        return V_list 

    return V

policy_evaluation

policy_evaluation(env, discount, policy, precision=0.001, max_backups=1000)

Evaluates a given policy to compute its value function.

Parameters:

Name Type Description Default
env

A gym-classics environment with model access.

required
discount

The discount factor (0 <= discount <= 1).

required
policy

The policy to evaluate.

required
precision

The precision for convergence (default: 1e-3).

0.001
max_backups

Maximum number of backups to perform to prevent infinite loops (default: 1000).

1000

Returns:

Type Description

The value function for the given policy.

Source code in gym_classics2/algorithms/dynamic_programming.py
def policy_evaluation(env, discount, policy, precision=1e-3, max_backups=1000):
    """Evaluates a given policy to compute its value function.

    Args:
        env: A gym-classics environment with model access.
        discount: The discount factor (0 <= discount <= 1).
        policy: The policy to evaluate.
        precision: The precision for convergence (default: 1e-3).
        max_backups: Maximum number of backups to perform to prevent infinite loops (default: 1000).

    Returns:
        The value function for the given policy.
    """

    assert isinstance(env, GymClassicsBaseEnv), "Value iteration requires a gym-classics environment with model access to the environment." 
    assert 0.0 <= discount <= 1.0
    assert precision > 0.0

    V = np.zeros(len(policy), dtype=np.float64)

    while True:
        V_old = V.copy()

        for s in env.states():
            V[s] = backup(env, discount, V, s, policy[s])

        if np.abs(V - V_old).max() <= precision or max_backups <= 0:
            break

        max_backups -= 1
    return V

policy_improvement

policy_improvement(env, discount, policy, V_policy, precision=0.001)

Improves the policy based on the given value function.

Parameters:

Name Type Description Default
env

A gym-classics environment with model access.

required
discount

The discount factor (0 <= discount <= 1).

required
policy

The current policy to improve.

required
V_policy

The value function of the current policy.

required
precision

The precision for determining stability (default: 1e-3).

0.001

Returns:

Type Description

A tuple (improved_policy, stable) where stable is True if the policy did not change.

Source code in gym_classics2/algorithms/dynamic_programming.py
def policy_improvement(env, discount, policy, V_policy, precision=1e-3):
    """Improves the policy based on the given value function.

    Args:
        env: A gym-classics environment with model access.
        discount: The discount factor (0 <= discount <= 1).
        policy: The current policy to improve.
        V_policy: The value function of the current policy.
        precision: The precision for determining stability (default: 1e-3).

    Returns:
        A tuple (improved_policy, stable) where stable is True if the policy did not change.
    """

    policy_old = policy.copy()
    V_old = V_policy.copy()

    for s in env.states():
        Q_values = [backup(env, discount, V_policy, s, a) for a in env.actions()]
        policy[s] = np.argmax(Q_values)
        V_policy[s] = max(Q_values)

    stable = np.logical_or(
        policy == policy_old,
        np.abs(V_policy - V_old).max() <= precision,
    ).all()

    return policy, stable

policy_iteration

policy_iteration(env, discount, precision=0.001, max_backups=1000, history=False, verbose=False, rng=None)

Performs policy iteration for the given environment.

Parameters:

Name Type Description Default
env

A gym-classics environment with model access.

required
discount

The discount factor (0 <= discount <= 1).

required
precision

The precision for convergence (default: 1e-3).

0.001
max_backups

Maximum number of iterations used in policy evaluation. Note: this prevents an infinite loop for policies that do not reach a terminal state.

1000
history

If True, returns lists of intermediate policies and value functions.

False
verbose

If True, prints progress information.

False
rng

NumPy generator or integer seed used to initialize the policy.

None

Returns:

Type Description

The optimal policy. If history is True, returns a tuple (policy_list, V_list) containing lists of intermediate policies and value functions.

Source code in gym_classics2/algorithms/dynamic_programming.py
def policy_iteration(env, discount, precision=1e-3, max_backups=1000, history=False, verbose=False, rng=None):
    """Performs policy iteration for the given environment.

    Args:
        env: A gym-classics environment with model access.
        discount: The discount factor (0 <= discount <= 1).
        precision: The precision for convergence (default: 1e-3).
        max_backups: Maximum number of iterations used in policy evaluation. Note: this prevents an infinite loop for policies that do not reach a terminal state.
        history: If True, returns lists of intermediate policies and value functions.
        verbose: If True, prints progress information.
        rng: NumPy generator or integer seed used to initialize the policy.

    Returns:
        The optimal policy. If history is True, returns a tuple (policy_list, V_list) containing lists of intermediate policies and value functions.
    """

    assert isinstance(env, GymClassicsBaseEnv), "Value iteration requires a gym-classics environment with model access to the environment." 
    assert 0.0 <= discount <= 1.0
    assert precision > 0.0

    policy = random_policy(env, rng=rng)

    if history:
        pol_list = []
        pol_list.append(policy.copy())
        V_list = []

    iterations = 0
    progress = tqdm(total=None, desc="Policy Iteration", disable=verbose)
    while True:
        progress.update()
        if verbose:
            print('.', end = '')
            iterations += 1

        V_policy = policy_evaluation(env, discount, policy, precision, max_backups)
        if history:
            V_list.append(V_policy.copy())

        policy, stable = policy_improvement(env, discount, policy, V_policy, precision)

        if stable:
            break

        if history:
            pol_list.append(policy.copy())

    if verbose:
        print(f'\nConverged after {iterations} iterations.')

    if history:
        return pol_list, V_list

    return policy

Policy helpers

gym_classics2.algorithms.policy

This file implements different policy representations and functions for working with policies in gym-classics environments. It includes functions for creating random policies, encoding policies for display, and computing greedy policies based

make_multidiscrete_policy

make_multidiscrete_policy(policy, env)

Converts a tabular policy vector to a multi-discrete tabular policy stored in a dictionary that can be used for sample.

Source code in gym_classics2/algorithms/policy.py
def make_multidiscrete_policy(policy, env):
    """
    Converts a tabular policy vector to a multi-discrete tabular policy stored in a dictionary that can be used for sample.
    """
    assert isinstance(env.observation_space, gym.spaces.MultiDiscrete), "Requires an environment with a multi-discrete state spaces."

    if isinstance(policy, dict):
        return policy

    if isinstance(env, GymClassicsBaseEnv):
        return dict(zip(env.id2state(list(env.states())), [policy[s] for s in env.states()]))
    else:
        raise ValueError("Unsupported environment type")

random_policy

random_policy(env, rng=None)

Create a random policy for the given environment. The policy is represented as a numpy array where each entry corresponds to an action for a state. rng may be a NumPy generator or an integer seed.

Source code in gym_classics2/algorithms/policy.py
def random_policy(env, rng=None):
    """
    Create a random policy for the given environment. 
    The policy is represented as a numpy array where each entry corresponds to an action for a state.
    ``rng`` may be a NumPy generator or an integer seed.
    """

    rng = get_rng(rng)

    if isinstance(env.observation_space, gym.spaces.Discrete):
        return rng.integers(env.action_space.n, size=env.observation_space.n)
    else:
        assert isinstance(env, GymClassicsBaseEnv), "Only gym-classics environments are supported for random policies with multi-discrete state spaces."
        return make_multidiscrete_policy(rng.integers(env.action_space.n, size=len(env.states())), env)

encode_policy

encode_policy(env, policy, type='text')

Encode a policy for display. The policy is represented as a numpy array where each entry corresponds to an action for a state. The function returns a list of action names corresponding to the actions in the policy.

Source code in gym_classics2/algorithms/policy.py
def encode_policy(env, policy, type = "text"):
    """
    Encode a policy for display. The policy is represented as a numpy array where each entry corresponds to an action for a state.
    The function returns a list of action names corresponding to the actions in the policy.
    """

    assert isinstance(env, GymClassicsBaseEnv)

    return [env.unwrapped.id2action(a, type = type) for a in policy]

greedy_policy

greedy_policy(env, V, discount=1, rng=None)

Calculate the greedy policy for a given value function.

:param env: the environment :param V: the value function as a 1-D numpy array :param discount: discount factor :param rng: NumPy generator or integer seed used for random tie-breaking

Source code in gym_classics2/algorithms/policy.py
def greedy_policy(env, V, discount=1, rng=None):
    """
    Calculate the greedy policy for a given value function.

    :param env: the environment
    :param V: the value function as a 1-D numpy array
    :param discount: discount factor
    :param rng: NumPy generator or integer seed used for random tie-breaking
    """
    assert isinstance(env, GymClassicsBaseEnv), "greedy_policy requires a gym-classics environment with discrete state space."
    assert isinstance(env.action_space, gym.spaces.Discrete)
    assert isinstance(env.observation_space, gym.spaces.Discrete)
    assert 0.0 <= discount <= 1.0

    policy = np.zeros(len(env.states()), dtype=np.int64)
    rng = get_rng(rng)

    env = env.unwrapped

    for s in env.states():
        Q_values = [_backup(env, discount, V, s, a) for a in range(env.action_space.n)]
        policy[s] = random_argmax(Q_values, rng=rng)

    return policy

greedy_policy_Q

greedy_policy_Q(env, Q, discount=1, rng=None)

Calculate the greedy policy for a given value function.

:param env: the environment :param Q: the action value function as a state-by-action numpy array :param rng: NumPy generator or integer seed used for random tie-breaking :param discount: discount factor

Source code in gym_classics2/algorithms/policy.py
def greedy_policy_Q(env, Q, discount=1, rng=None):
    """
    Calculate the greedy policy for a given value function.

    :param env: the environment
    :param Q: the action value function as a state-by-action numpy array
    :param rng: NumPy generator or integer seed used for random tie-breaking
    :param discount: discount factor
    """

    assert isinstance(env.action_space, gym.spaces.Discrete)
    assert isinstance(env.observation_space, gym.spaces.Discrete)
    assert 0.0 <= discount <= 1.0

    return random_argmax(Q, axis=1, rng=rng)

epsilon_greedy_action

epsilon_greedy_action(policy, state=None, epsilon=0, rng=None)

Get an epsilon-greedy action for a given tabular policy.

:param policy: the policy as a 1-D numpy array :param epsilon: the probability of taking a random action :param state: the current state :param rng: NumPy generator or integer seed

Source code in gym_classics2/algorithms/policy.py
def epsilon_greedy_action(policy, state=None, epsilon=0, rng=None):
    """
    Get an epsilon-greedy action for a given tabular policy.

    :param policy: the policy as a 1-D numpy array
    :param epsilon: the probability of taking a random action
    :param state: the current state
    :param rng: NumPy generator or integer seed
    """

    rng = get_rng(rng)

    if epsilon > 0 and rng.random() < epsilon:
        return rng.integers(len(policy)) if state is None else rng.integers(len(policy[state]))

    if state is None:
        return random_argmax(policy, rng=rng)
    else:
        return random_argmax(policy[state], rng=rng)

Monte Carlo methods

gym_classics2.algorithms.monte_carlo_methods

This file implements tabular Monte Carlo methods for policy evaluation and control in gym-classics environments with discrete state spaces.

states

states(env)

Returns a list of all states in the environment.

Source code in gym_classics2/algorithms/monte_carlo_methods.py
def states(env):
    """Returns a list of all states in the environment."""
    assert isinstance(env.observation_space, gym.spaces.Discrete) or isinstance(env.observation_space, gym.spaces.MultiDiscrete), "Tabular methods require discrete state space."  

    if isinstance(env.observation_space, gym.spaces.Discrete):
        return list(range(env.observation_space.n))
    elif isinstance(env.observation_space, gym.spaces.MultiDiscrete):
        return [tuple(s) for s in np.array(np.meshgrid(*[range(n) for n in env.observation_space.nvec])).T.reshape(-1, len(env.observation_space.nvec))]
    else:
        raise ValueError("Unsupported observation space type for state enumeration.")

sample_episode

sample_episode(env, policy=None, start_state=None, start_action=None, epsilon=0, max_len=1000, verbose=False, rng=None)

Samples an episode from the environment using the given policy and starting conditions.

Parameters:

Name Type Description Default
env

The environment to sample from.

required
policy

A mapping from states to actions. For discrete state spaces, use an action list in state order. If None, a random policy is used.

None
start_state

The state to start the episode from (if None, the environment's default starting state will be used).

None
start_action

The action to take in the first step of the episode (if None, the action will be chosen according to the policy or randomly if no policy is given).

None
epsilon

The probability of taking a random action instead of the policy's action at each step (for epsilon-greedy exploration).

0
max_len

The maximum length of the episode to prevent infinite loops.

1000
verbose

If True, prints the state transitions and rewards for each step in the episode.

False
rng

NumPy generator or integer seed for policy and exploration choices.

None

Returns:

Type Description

A list of (state, action, reward, next_state) tuples representing the episode.

The episode ends when a terminal state is reached or when max_len steps have been taken.

Note that the last tuple in the episode will have a next_state that is either terminal or the

state at which the episode was truncated due to max_len.

Source code in gym_classics2/algorithms/monte_carlo_methods.py
def sample_episode(env, policy=None, start_state=None, start_action=None, epsilon=0,
                   max_len=1000, verbose=False, rng=None):
    """
    Samples an episode from the environment using the given policy and starting conditions.

    Args: 
        env: The environment to sample from.
        policy: A mapping from states to actions. For discrete state spaces, use
            an action list in state order. If ``None``, a random policy is used.
        start_state: The state to start the episode from (if None, the environment's default starting state will be used).
        start_action: The action to take in the first step of the episode (if None, the action will be chosen according to the policy or randomly if no policy is given).
        epsilon: The probability of taking a random action instead of the policy's action at each step (for epsilon-greedy exploration).
        max_len: The maximum length of the episode to prevent infinite loops.
        verbose: If True, prints the state transitions and rewards for each step in the episode.
        rng: NumPy generator or integer seed for policy and exploration choices.

    Returns:
        A list of (state, action, reward, next_state) tuples representing the episode. 
        The episode ends when a terminal state is reached or when max_len steps have been taken.    
        Note that the last tuple in the episode will have a next_state that is either terminal or the 
        state at which the episode was truncated due to max_len.  
    """

    assert isinstance(env.observation_space, gym.spaces.Discrete) or isinstance(env.observation_space, gym.spaces.MultiDiscrete), "Tabular methods require discrete state space."  
    assert 0.0 <= epsilon <= 1.0
    assert max_len > 0

    rng = get_rng(rng)

    if policy is None:
        policy = random_policy(env, rng=rng)
        epsilon = 1.0
        if verbose:
            print("*** No policy given, sampling using random actions!")

    if isinstance(env.observation_space, gym.spaces.MultiDiscrete):
        policy = make_multidiscrete_policy(policy, env)

    episode = []
    s, r = env.reset()

    # force custom start state could be a state or an state index.
    if not start_state is None:
        if isinstance(env, BaseEnv) and env.tabular:
            env.state = env.id2state(start_state)
        elif isinstance(env.observation_space, gym.spaces.Discrete):
            env.state = int(start_state)
        else:
            raise ValueError("Custom start state is only supported for tabular gym-classics environments and environments with a MultiDiscrete observation space.")
        s = start_state

    step = 0
    done = False
    while not done and step < max_len:
        a = policy[s]

        # epsilon greedy choice?
        if epsilon > 0 and rng.random() < epsilon:
            a = rng.integers(env.action_space.n)

        # for exploring starts
        if step == 0 and not start_action is None:
            a = start_action

        sp, reward, done, _, _ = env.step(a)

        episode.append((s,a,reward,sp))

        if verbose:
            print ("Step", step, "(s,a,r,s')=",s,a,reward,sp)

        s = sp
        step += 1

    return(episode)

on_policy_state_distribution

on_policy_state_distribution(env, pol, discount=1, epsilon=0, n=100, verbose=False, rng=None)

Estimate a policy's state distribution using rng for exploration.

Source code in gym_classics2/algorithms/monte_carlo_methods.py
def on_policy_state_distribution(env, pol, discount=1, epsilon=0, n=100,
                                 verbose=False, rng=None):
    """Estimate a policy's state distribution using ``rng`` for exploration."""

    assert isinstance(env.observation_space, gym.spaces.Discrete), "Tabular methods require discrete state space."   
    assert 0.0 <= epsilon <= 1.0
    assert 0.0 < discount <= 1.0
    assert n > 0

    rng = get_rng(rng)
    state_counts = np.zeros(env.observation_space.n)

    for _ in tqdm(range(n), desc="Sampling Episodes", disable=verbose):
        episode = np.array(sample_episode(env, policy=pol, epsilon=epsilon, rng=rng))
        states = np.append(episode[:,0], episode[-1,3])
        discounts = np.array([discount**i for i in range(len(states))])
        for s, d in zip(states, discounts):
            state_counts[int(s)] += d

    state_prob = state_counts / sum(state_counts)
    return state_prob

MC_prediction

MC_prediction(env, policy, discount, n=100, max_episode_len=100, verbose=False, rng=None)

Estimate a policy's state values with first-visit Monte Carlo prediction.

Parameters:

Name Type Description Default
env

Gymnasium environment with a discrete observation space.

required
policy

Array-like mapping from state IDs to action IDs.

required
discount

Reward discount factor.

required
n

Number of episodes to sample.

100
max_episode_len

Maximum sampled steps per episode.

100
verbose

Print episode progress when true.

False
rng

NumPy generator or integer seed for episode sampling.

None

Returns:

Type Description

A NumPy value array with one entry per state. Unvisited states contain

numpy.nan.

Source code in gym_classics2/algorithms/monte_carlo_methods.py
def MC_prediction(env, policy, discount, n=100, max_episode_len=100,
                  verbose=False, rng=None):
    """Estimate a policy's state values with first-visit Monte Carlo prediction.

    Args:
        env: Gymnasium environment with a discrete observation space.
        policy: Array-like mapping from state IDs to action IDs.
        discount: Reward discount factor.
        n: Number of episodes to sample.
        max_episode_len: Maximum sampled steps per episode.
        verbose: Print episode progress when true.
        rng: NumPy generator or integer seed for episode sampling.

    Returns:
        A NumPy value array with one entry per state. Unvisited states contain
        ``numpy.nan``.
    """
    assert isinstance(env.observation_space, gym.spaces.Discrete), "Tabular methods require discrete state space."  
    assert n > 0
    assert max_episode_len > 0

    rng = get_rng(rng)
    Returns = defaultdict(list) # a list for each s

    for i in tqdm(range(n), desc="MC Prediction", disable=verbose):
        if verbose:
            print("episode", i," of ", n)

        episode = sample_episode(env, policy, max_len=max_episode_len, rng=rng)
        G = 0

        # for first visit check for s
        visited = [s for s,a,r,sp in episode]

        # process episode in reverse order
        for t in range(len(episode)-1, -1, -1):
            s,a,r,sp = episode[t]
            G = discount * G + r

            # use first visit of s only
            if not s in visited[0:t]:
                Returns[s].append(G)

    Vs = np.array([np.mean(Returns[s]) if len(Returns[s]) else np.nan for s in range(env.observation_space.n)])
    return Vs

MC_control_ES_textbook

MC_control_ES_textbook(env, discount, n=100, Q=None, max_episode_len=100, history=False, verbose=False, rng=None)

Monte Carlo control with exploring starts and stored sample returns.

This direct textbook implementation stores every first-visit return. Prefer :func:MC_control_ES for larger experiments because its running-average update uses less memory.

Parameters:

Name Type Description Default
env

A tabular gym_classics2 environment.

required
discount

Reward discount factor.

required
n

Number of episodes to sample.

100
Q

Optional initial action-value array.

None
max_episode_len

Maximum sampled steps per episode.

100
history

Retain intermediate policies, Q arrays, episodes, and returns.

False
verbose

Print episode details; values greater than one print transitions.

False
rng

NumPy generator or integer seed for all algorithm choices.

None

Returns:

Type Description

(policy, Q). If history=True, a third item contains the history

dictionary with policies, Q_values, episodes, and returns.

Source code in gym_classics2/algorithms/monte_carlo_methods.py
def MC_control_ES_textbook(env, discount, n=100, Q=None, max_episode_len=100,
                           history=False, verbose=False, rng=None):
    """Monte Carlo control with exploring starts and stored sample returns.

    This direct textbook implementation stores every first-visit return. Prefer
    :func:`MC_control_ES` for larger experiments because its running-average
    update uses less memory.

    Args:
        env: A tabular ``gym_classics2`` environment.
        discount: Reward discount factor.
        n: Number of episodes to sample.
        Q: Optional initial action-value array.
        max_episode_len: Maximum sampled steps per episode.
        history: Retain intermediate policies, Q arrays, episodes, and returns.
        verbose: Print episode details; values greater than one print transitions.
        rng: NumPy generator or integer seed for all algorithm choices.

    Returns:
        ``(policy, Q)``. If ``history=True``, a third item contains the history
        dictionary with ``policies``, ``Q_values``, ``episodes``, and ``returns``.
    """
    assert isinstance(env.observation_space, gym.spaces.Discrete), "Tabular methods require discrete state space."  
    assert n > 0
    assert max_episode_len > 0
    assert Q is None or Q.shape == (len(env.states()), env.action_space.n)

    rng = get_rng(rng)
    policy = random_policy(env, rng=rng)

    if Q is None:
        Q = np.zeros((len(env.states()), env.action_space.n))

    # lists that grow are very slow. We should use a running average instead. But this is easier to read and understand.
    Returns = defaultdict(list) # a list for each (s,a)

    if history:
        Q_list = []
        Q_list.append(Q.copy())
        pol_list = []
        pol_list.append(policy.copy())
        ep_list = []
        return_list = []


    for i in tqdm(range(n), desc="MC Control", disable=verbose):
        if verbose:
            print("episode", i," of ", n)

        # Sample starting (s,a)
        s = rng.integers(env.observation_space.n)
        a = rng.integers(env.action_space.n)

        episode = sample_episode(env, policy, start_state=s, start_action=a,
                                 max_len=max_episode_len, verbose=verbose > 1,
                                 rng=rng)
        G = 0

        # for first visit check for (s,a)
        visited = [(s,a) for s,a,r,sp in episode]

        # process episode in reverse order
        for t in range(len(episode)-1, -1, -1):
            s,a,r,sp = episode[t]
            G = discount * G + r

            # use first visit of (s,a) only
            if not (s,a) in visited[0:t]:
                Returns[(s,a)].append(G)

                # update policy
                Q[s,a] = np.mean(Returns[(s,a)])
                policy[s] = random_argmax(Q[s, :], rng=rng)

        if history:
            pol_list.append(policy.copy())
            Q_list.append(Q.copy())
            ep_list.append(episode.copy())
            return_list.append(G)

    if history:
        return policy, Q, {'policies': pol_list, 'Q_values': Q_list, 'episodes': ep_list, 'returns': return_list}

    return policy, Q

MC_control_ES

MC_control_ES(env, discount, n=100, Q=None, max_episode_len=100, history=False, verbose=False, rng=None)

Monte Carlo Control with Exploring Starts (incremental version). This algorithm estimates the optimal action-value function Q and the corresponding greedy policy by sampling episodes with exploring starts. It uses incremental updates to compute the average returns for each (s,a) pair, which is more memory efficient than storing all returns. Args: env: The environment to interact with. Must have discrete state and action spaces. discount: The discount factor (gamma) for future rewards. Should be in (0, 1]. n: The number of episodes to sample for learning. Must be a positive integer. Q: Optional initial action-value function. If None, it will be initialized to zeros. max_episode_len: Maximum length of each episode to prevent infinite loops. Must be a positive integer. history: If True, the function will return the history of policies, Q-values, and and episodes for each iteration. This can be useful for analysis and visualization, but it will consume more memory. verbose: If True, the function will print progress and episode details. If verbose > 1, it will also print the state transitions and rewards for each step in the episode. rng: NumPy generator or integer seed for all algorithm choices. Returns: If history is False: A tuple (policy, Q) where policy is the learned greedy policy and Q is the learned action-value function. If history is True: A tuple (pol_list, Q_list, ep_list) where pol_list is a list of policies for each iteration, Q
is a list of Q-value functions for each iteration, and ep_list is a list of episodes sampled in each iteration.

Source code in gym_classics2/algorithms/monte_carlo_methods.py
def MC_control_ES(env, discount, n=100, Q=None, max_episode_len=100,
                  history=False, verbose=False, rng=None):
    """Monte Carlo Control with Exploring Starts (incremental version).
    This algorithm estimates the optimal action-value function Q and the corresponding greedy policy by sampling episodes with exploring starts. It uses incremental updates to compute the average returns for each (s,a) pair, which is more memory efficient than storing all returns.
    Args: env: The environment to interact with. Must have discrete state and action spaces.
        discount: The discount factor (gamma) for future rewards. Should be in (0, 1].
        n: The number of episodes to sample for learning. Must be a positive integer.
        Q: Optional initial action-value function. If None, it will be initialized to zeros.
        max_episode_len: Maximum length of each episode to prevent infinite loops. Must be a positive integer.
        history: If True, the function will return the history of policies, Q-values, and
                 and episodes for each iteration. This can be useful for analysis and visualization, but it will consume more memory.
        verbose: If True, the function will print progress and episode details. If verbose > 1, it will also print the state transitions and rewards for each step in the episode.
        rng: NumPy generator or integer seed for all algorithm choices.
    Returns:    If history is False: A tuple (policy, Q) where policy is the learned greedy policy and Q is the learned action-value function.
        If history is True: A tuple (pol_list, Q_list, ep_list) where pol_list is a list of policies for each iteration, Q          
        is a list of Q-value functions for each iteration, and ep_list is a list of episodes sampled in each iteration.
    """

    assert isinstance(env.observation_space, gym.spaces.Discrete), "Tabular methods require discrete state space."  
    assert n > 0
    assert max_episode_len > 0
    assert Q is None or Q.shape == (env.observation_space.n, env.action_space.n)

    rng = get_rng(rng)
    policy = random_policy(env, rng=rng)

    if Q is None:
        Q = np.zeros((env.observation_space.n, env.action_space.n))

    # Count how many first-visit returns have been used for each (s, a)
    N = np.zeros((env.observation_space.n, env.action_space.n), dtype=int)

    if history:
        Q_list = []
        Q_list.append(Q.copy())
        pol_list = []
        pol_list.append(policy.copy())
        ep_list = []
        return_list = []

    for i in tqdm(range(n), desc="MC Control (Incremental)", disable=verbose):
        if verbose:
            print("episode", i, "of", n)

        # Exploring starts
        s = rng.integers(env.observation_space.n)
        a = rng.integers(env.action_space.n)

        episode = sample_episode(
            env,
            policy,
            start_state=s,
            start_action=a,
            max_len=max_episode_len,
            verbose=verbose > 1,
            rng=rng,
        )

        G = 0

        # Process episode backward
        visited = set()
        for t in range(len(episode) - 1, -1, -1):
            s, a, r, sp = episode[t]
            G = discount * G + r

            # First-visit MC: only update the first occurrence from the start
            if (s, a) not in visited:
                visited.add((s, a))

                N[s, a] += 1
                Q[s, a] += (G - Q[s, a]) / N[s, a]

                # Improve policy greedily
                policy[s] = random_argmax(Q[s, :], rng=rng)

        if history:
            pol_list.append(policy.copy())
            Q_list.append(Q.copy())
            ep_list.append(episode.copy())
            return_list.append(G)

    if history:
        return policy, Q, {'policies': pol_list, 'Q_values': Q_list, 'episodes': ep_list, 'returns': return_list}

    return policy, Q

Temporal-difference learning

gym_classics2.algorithms.temporal_difference_learning

This file implements temporal difference learning algorithms for policy evaluation and control in gym-classics
environments with discrete state spaces. The algorithms include Sarsa(0) and Q-learning, which are fundamental methods in reinforcement learning for learning value functions and optimal policies from experience without requiring a model of the environment.

Sarsa_0

Sarsa_0(env, discount, alpha, epsilon, Q=None, n=100, verbose=False, history=False, rng=None)

Learn action values with one-step on-policy Sarsa.

Parameters:

Name Type Description Default
env

Gymnasium environment with discrete observation and action spaces.

required
discount

Reward discount factor in [0, 1].

required
alpha

Scalar step size or a :class:~gym_classics2.algorithms.schedules.Schedule evaluated once per episode.

required
epsilon

Scalar exploration probability or a schedule evaluated once per episode.

required
Q

Optional initial action-value array shaped (env.observation_space.n, env.action_space.n). The array is updated in place.

None
n

Number of training episodes.

100
verbose

Print individual updates when true.

False
history

Retain Q arrays, discounted episode returns, and episode lengths.

False
rng

NumPy generator or integer seed for exploration and tie-breaking.

None

Returns:

Type Description

The learned Q array. If history=True, returns (Q, history_dict);

the dictionary contains Qs, returns, and ep_lens.

Raises:

Type Description
AssertionError

If the observation space is not discrete.

Source code in gym_classics2/algorithms/temporal_difference_learning.py
def Sarsa_0(env, discount, alpha, epsilon, Q=None, n=100, verbose=False,
            history=False, rng=None):
    """Learn action values with one-step on-policy Sarsa.

    Args:
        env: Gymnasium environment with discrete observation and action spaces.
        discount: Reward discount factor in ``[0, 1]``.
        alpha: Scalar step size or a
            :class:`~gym_classics2.algorithms.schedules.Schedule` evaluated once
            per episode.
        epsilon: Scalar exploration probability or a schedule evaluated once per
            episode.
        Q: Optional initial action-value array shaped
            ``(env.observation_space.n, env.action_space.n)``. The array is updated
            in place.
        n: Number of training episodes.
        verbose: Print individual updates when true.
        history: Retain Q arrays, discounted episode returns, and episode lengths.
        rng: NumPy generator or integer seed for exploration and tie-breaking.

    Returns:
        The learned Q array. If ``history=True``, returns ``(Q, history_dict)``;
        the dictionary contains ``Qs``, ``returns``, and ``ep_lens``.

    Raises:
        AssertionError: If the observation space is not discrete.
    """
    assert isinstance(env.observation_space, gym.spaces.Discrete), "Tabular methods require discrete state space."  

    rng = get_rng(rng)

    if not isinstance(alpha, Schedule):
        alpha = ConstantSchedule(alpha)
    if not isinstance(epsilon, Schedule):
        epsilon = ConstantSchedule(epsilon)

    if Q is None:
        Q = np.zeros((env.observation_space.n, env.action_space.n))

    if history:
        Q_list = []
        Q_list.append(Q.copy())
        return_list = []
        ep_len_list = []


    for i in tqdm(range(n), desc="Sarsa", disable=verbose):
        s, r = env.reset()

        if verbose:
            print(f"--- Episode {i} ---")      

        if rng.random() > epsilon(i):
            a = random_argmax(Q[s, :], rng=rng)
        else:
            a = rng.integers(env.action_space.n)

        t = 0
        done = False
        G = 0
        while not done:            
            sp, r, done, _, _ = env.step(a)
            G += r * pow(discount, t)
            t += 1


            if rng.random() > epsilon(i):
                ap = random_argmax(Q[sp, :], rng=rng)
            else:
                ap = rng.integers(env.action_space.n)

            Q[s,a] = Q[s,a] + alpha(i) * (r + discount * Q[sp,ap] - Q[s,a])

            if verbose:
                print(f"{t} - sarsa: {s},{a},{r},{sp},{ap}, - new Q(s,a): {Q[s,a]}")
                if done:
                    print("Total return:", G)

            s = sp
            a = ap

        if history:
            Q_list.append(Q.copy())
            return_list.append(G)   
            ep_len_list.append(t)

    if history:
        return Q, {'Qs': Q_list, 'returns': return_list, 'ep_lens': ep_len_list}

    return Q

Q_learning

Q_learning(env, discount, alpha, epsilon, Q=None, n=100, verbose=False, history=False, rng=None)

Learn action values with one-step off-policy Q-learning.

Parameters:

Name Type Description Default
env

Gymnasium environment with discrete observation and action spaces.

required
discount

Reward discount factor in [0, 1].

required
alpha

Scalar step size or a :class:~gym_classics2.algorithms.schedules.Schedule evaluated once per episode.

required
epsilon

Scalar exploration probability or a schedule evaluated once per episode.

required
Q

Optional initial action-value array shaped (env.observation_space.n, env.action_space.n). The array is updated in place.

None
n

Number of training episodes.

100
verbose

Print progress information when true.

False
history

Retain Q arrays, discounted episode returns, episode lengths, and state-visit counts.

False
rng

NumPy generator or integer seed for exploration and tie-breaking.

None

Returns:

Type Description

The learned Q array. If history=True, returns (Q, history_dict);

the dictionary contains Qs, returns, ep_lens, and

state_visits.

Raises:

Type Description
AssertionError

If the observation space is not discrete.

Source code in gym_classics2/algorithms/temporal_difference_learning.py
def Q_learning(env, discount, alpha, epsilon, Q=None, n=100, verbose=False,
               history=False, rng=None):
    """Learn action values with one-step off-policy Q-learning.

    Args:
        env: Gymnasium environment with discrete observation and action spaces.
        discount: Reward discount factor in ``[0, 1]``.
        alpha: Scalar step size or a
            :class:`~gym_classics2.algorithms.schedules.Schedule` evaluated once
            per episode.
        epsilon: Scalar exploration probability or a schedule evaluated once per
            episode.
        Q: Optional initial action-value array shaped
            ``(env.observation_space.n, env.action_space.n)``. The array is updated
            in place.
        n: Number of training episodes.
        verbose: Print progress information when true.
        history: Retain Q arrays, discounted episode returns, episode lengths, and
            state-visit counts.
        rng: NumPy generator or integer seed for exploration and tie-breaking.

    Returns:
        The learned Q array. If ``history=True``, returns ``(Q, history_dict)``;
        the dictionary contains ``Qs``, ``returns``, ``ep_lens``, and
        ``state_visits``.

    Raises:
        AssertionError: If the observation space is not discrete.
    """
    assert isinstance(env.observation_space, gym.spaces.Discrete), "Tabular methods require discrete state space."  

    rng = get_rng(rng)

    if not isinstance(alpha, Schedule):
        alpha = ConstantSchedule(alpha)
    if not isinstance(epsilon, Schedule):
        epsilon = ConstantSchedule(epsilon)

    if Q is None:
        Q = np.zeros((env.observation_space.n, env.action_space.n))

    if history:
        Q_list = []
        Q_list.append(Q.copy())
        return_list = []
        ep_len_list = []
        #state_visits = np.zeros(env.observation_space.n, dtype=int)
        # Note we use float so visualization works better
        state_visits = np.zeros(env.observation_space.n, dtype=float)

    for i in tqdm(range(n), desc="Q-Learning", disable=verbose):
        s, r = env.reset()

        if history:
            state_visits[s] += 1

        done = False
        G = 0
        t = 0   
        while not done:
            # epsilon-greedy choice w.r.t. Q 
            if rng.random() > epsilon(i):
                a = random_argmax(Q[s, :], rng=rng)
            else:
                a = rng.integers(env.action_space.n)

            sp, r, done, _, _ = env.step(a)

            if history:
                state_visits[sp] += 1

            if history:
                G += r * pow(discount, t)
                t += 1

            Q[s,a] = Q[s,a] + alpha(i) * (r + discount * np.max(Q[sp,:]) - Q[s,a])

            s = sp

        if history:
            Q_list.append(Q.copy())
            return_list.append(G)
            ep_len_list.append(t)

    if history:
        return Q, {'Qs': Q_list, 'returns': return_list, 'ep_lens': ep_len_list, 'state_visits': state_visits}

    return Q

Linear function approximation

gym_classics2.algorithms.linear_approximation

This file implements linear function approximation algorithms for policy evaluation and control. This is not a tabular approach and does not require discrete state spaces. The user needs to implement the state_features function to convert states to feature vectors.

state_features

state_features(s, env)

Converts a state to a state feature vector. It needs to be overwritten by the user to implement different feature representations. This could be linear features, tile coding, radial basis functions, Fourier basis functions, or even a neural network.

:param s: state :param env: environment instance

:return a state feature vector

Source code in gym_classics2/algorithms/linear_approximation.py
def state_features(s,env):
    """
    Converts a state to a state feature vector. It needs to be overwritten by the user to implement different feature representations. 
    This could be linear features, tile coding, radial basis functions, Fourier basis functions, or even a neural network.

    :param s: state
    :param env: environment instance   

    :return a state feature vector
    """
    raise NotImplementedError("state_features function needs to be implemented by the user. By default, it just concatenates a constant feature (for the intercept) with the state itself. This is equivalent to linear function approximation with a tabular representation.")

active_weights

active_weights(a, sf_len)

helper for q_hat()

Source code in gym_classics2/algorithms/linear_approximation.py
def active_weights(a, sf_len):
    """helper for q_hat()"""
    return [0] + list(range(a*sf_len+1, a*sf_len+sf_len+1))

state_action_features

state_action_features(s, a, env)

Construct a block-coded feature vector for a state-action pair.

Source code in gym_classics2/algorithms/linear_approximation.py
def state_action_features(s,a,env):
    """Construct a block-coded feature vector for a state-action pair."""
    s = state_features(s,env)
    x = np.zeros(1+len(s)*env.action_space.n)
    x[active_weights(a, len(s)-1)] = s
    return x

v_hat

v_hat(s, w, env)

Estimate Value function

:param s: state id :param w: weight vector :param env: environment instance

:return the state value estimate

Source code in gym_classics2/algorithms/linear_approximation.py
def v_hat(s, w, env):
    """
    Estimate Value function

    :param s: state id
    :param w: weight vector
    :param env: environment instance

    :return the state value estimate
    """
    return np.dot(w, state_features(s, env))

q_hat

q_hat(s, a, w, env)

Estimate the action value function.

:param s: state id :param a: action :param w: weight vector :param env: environment instance :return the state-action value estimate

Source code in gym_classics2/algorithms/linear_approximation.py
def q_hat(s, a, w, env):
    """
    Estimate the action value function.

    :param s: state id
    :param a: action
    :param w: weight vector
    :param env: environment instance
    :return the state-action value estimate
    """    
    x = state_action_features(s, a, env)
    return np.dot(w, x)

epsilon_greedy_action_w

epsilon_greedy_action_w(env, w, state, epsilon=0, rng=None)

Get an epsilon-greedy action for a given policy.

:param w: weight vector for the action-value function approximator :param env: environment instance :param state: the current state :param epsilon: the probability of taking a random action :param rng: NumPy generator or integer seed

Source code in gym_classics2/algorithms/linear_approximation.py
def epsilon_greedy_action_w(env, w, state, epsilon=0, rng=None):
    """
    Get an epsilon-greedy action for a given policy.

    :param w: weight vector for the action-value function approximator
    :param env: environment instance
    :param state: the current state
    :param epsilon: the probability of taking a random action
    :param rng: NumPy generator or integer seed
    """

    rng = get_rng(rng)

    if epsilon > 0 and rng.random() < epsilon:
        return rng.integers(env.action_space.n)

    return random_argmax(
        [q_hat(state, a, w, env) for a in range(env.action_space.n)],
        rng=rng,
    )

MSVE

MSVE(V, V_true, weight=None)

Calculate the (weighted) mean squared value error.

:param V: value function to evaluate :param V_true: the value function to compare to :param weight: weight for each state. Typically the stationary state visit distribution.

Source code in gym_classics2/algorithms/linear_approximation.py
def MSVE(V, V_true, weight=None):
    """
    Calculate the (weighted) mean squared value error.

    :param V: value function to evaluate
    :param V_true: the value function to compare to
    :param weight: weight for each state. Typically the stationary state visit distribution.
    """
    if weight is None:
        weight = np.ones(len(V))

    return np.sum(weight * (V - V_true)**2)

semi_gradient_TD0_estimation

semi_gradient_TD0_estimation(env, policy, n, alpha, gamma, max_episode_length=1000, verbose=False)

Estimate the state-value function using the semi-gradient TD(0) algorithm.

This function runs TD(0) learning with function approximation over multiple episodes generated from a given policy and environment. Updates are performed using the semi-gradient of the value function approximation.

Parameters

env : Environment following the Gym interface from which episodes are sampled. policy : a deterministic policy as a vector. n : int Number of episodes to run for value estimation. Must be positive. alpha : float Step-size (learning rate) for TD updates. Must be in the interval (0, 1]. gamma : float Discount factor for future rewards. Must be in the interval [0, 1]. max_episode_length : int, optional Maximum number of time steps per episode (default is 1000). verbose : bool, optional If True, prints progress or diagnostic information during training (default is True).

Returns

w Returns the learned weight vector for the approximate value function.

Source code in gym_classics2/algorithms/linear_approximation.py
def semi_gradient_TD0_estimation(env, policy, n, alpha, gamma, max_episode_length=1000, verbose =False):
    """
    Estimate the state-value function using the semi-gradient TD(0) algorithm.

    This function runs TD(0) learning with function approximation over multiple
    episodes generated from a given policy and environment. Updates are performed
    using the semi-gradient of the value function approximation.

    Parameters
    ----------
    env : Environment following the Gym interface from which episodes are sampled.
    policy : a deterministic policy as a vector.
    n : int
        Number of episodes to run for value estimation. Must be positive.
    alpha : float
        Step-size (learning rate) for TD updates. Must be in the interval (0, 1].
    gamma : float
        Discount factor for future rewards. Must be in the interval [0, 1].
    max_episode_length : int, optional
        Maximum number of time steps per episode (default is 1000).
    verbose : bool, optional
        If True, prints progress or diagnostic information during training
        (default is True).

    Returns
    -------
    w
        Returns the learned weight vector for the approximate value function.
    """
    assert gamma >= 0 and gamma <= 1, "Gamma must be in [0,1]"
    assert n > 0, "Number of episodes must be positive"
    assert max_episode_length > 0, "Max episode length must be positive"

    if isinstance(env.observation_space, gym.spaces.Discrete):
        warnings.warn("The environment has a discrete state space. Consider using a tabular method instead of function approximation.")

    if not isinstance(alpha, Schedule):
        alpha = ConstantSchedule(alpha)

    state, _ = env.reset()
    w = np.zeros(len(state_features(state, env)))  # Initialize weights (intercept + x and y)

    for episode in tqdm(range(n), desc="Semi-Gradient TD(0)", disable=verbose):
        state, _ = env.reset()
        done = False

        i = 0
        while not done and i < max_episode_length:
            action = policy[state]  # follow policy
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            # Semi-gradient TD(0) update
            # Note: v_hat(terminal, w) needs to be 0
            if terminated:
                w += alpha(episode) * (reward - v_hat(state, w, env)) * state_features(state, env)    
            else: 
                w += alpha(episode) * (reward + gamma * v_hat(next_state, w, env) - v_hat(state, w, env)) * state_features(state, env)

            if verbose:
                print (f"Episode {episode+1}, Step {i+1}: S={state}, A={action}, R={reward}, S'={next_state}, w={w}")

            state = next_state
            i += 1

    return w

semi_gradient_Sarsa_0

semi_gradient_Sarsa_0(env, n, epsilon, alpha, gamma, w=None, max_episode_length=1000, verbose=False, history=False, rng=None)

Semi-gradient Sarsa(0): on-policy control with function approximation.

Implements the semi-gradient Sarsa(0) algorithm for estimating the optimal action-value function q_*(s, a) using a differentiable function approximator q̂(s, a, w). Actions are selected according to an ε-greedy policy derived from the current action-value estimate.

Episodes are truncated after max_episode_length time steps.

Parameters

env : Episodic environment used to generate experience. n : int Number of episodes over which to perform control learning. epsilon : float Exploration parameter for the epsilon-greedy behavior policy (0 <= epsilon <= 1). alpha : float Step-size parameter for the weight update (0 < alpha <= 1). gamma : float Discount factor (0 <= gamma <= 1). w : array-like or None, optional Initial weight vector for the action-value function approximator. If None, weights are initialized internally. max_episode_length : int, optional Maximum number of time steps per episode before truncation (default 1000). verbose : bool, optional If True, prints progress diagnostics during learning (default True). rng : numpy.random.Generator or int or None, optional Random generator or seed for exploration and tie-breaking.

Returns

w Returns the learned weight vector for the approximate value function.

Source code in gym_classics2/algorithms/linear_approximation.py
def semi_gradient_Sarsa_0(env, n, epsilon, alpha, gamma, w=None,
                          max_episode_length=1000, verbose=False,
                          history=False, rng=None):
    """
    Semi-gradient Sarsa(0): on-policy control with function approximation.

    Implements the **semi-gradient Sarsa(0)** algorithm for estimating the optimal
    action-value function q_*(s, a) using a differentiable function approximator
    q̂(s, a, w). Actions are selected according to an ε-greedy policy derived
    from the current action-value estimate.

    Episodes are truncated after `max_episode_length` time steps.

    Parameters
    ----------
    env : Episodic environment used to generate experience.
    n : int
        Number of episodes over which to perform control learning.
    epsilon : float
        Exploration parameter for the epsilon-greedy behavior policy (0 <= epsilon <= 1).
    alpha : float
        Step-size parameter for the weight update (0 < alpha <= 1).
    gamma : float
        Discount factor (0 <= gamma <= 1).
    w : array-like or None, optional
        Initial weight vector for the action-value function approximator.
        If None, weights are initialized internally.
    max_episode_length : int, optional
        Maximum number of time steps per episode before truncation (default 1000).
    verbose : bool, optional
        If True, prints progress diagnostics during learning (default True).
    rng : numpy.random.Generator or int or None, optional
        Random generator or seed for exploration and tie-breaking.

    Returns
    -------
    w
        Returns the learned weight vector for the approximate value function.
    """

    assert gamma >= 0 and gamma <= 1, "Gamma must be in [0,1]"
    assert n > 0, "Number of episodes must be positive"
    assert max_episode_length > 0, "Max episode length must be positive"

    rng = get_rng(rng)

    if isinstance(env.observation_space, gym.spaces.Discrete):
        warnings.warn("The environment has a discrete state space. Consider using a tabular method instead of function approximation.")

    if not isinstance(alpha, Schedule):
        alpha = ConstantSchedule(alpha)
    if not isinstance(epsilon, Schedule):
        epsilon = ConstantSchedule(epsilon)

    if w is None:
        state, _ = env.reset()
        w = np.zeros(len(state_action_features(state, 0, env)))

    if history:
        ws = []
        ws.append(w.copy())
        returns = []
        ep_lens = []

    for episode in tqdm(range(n), desc="Semi-Gradient SARSA(0)", disable=verbose):
        state, _ = env.reset()
        action = epsilon_greedy_action_w(env, w, state, epsilon(episode), rng=rng)
        done = False

        i = 0
        if history:
            G = 0
        while not done and i < max_episode_length:


            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            x = state_action_features(state, action, env)

            if terminated:
                next_action = None
                w += alpha(episode) * (reward - q_hat(state, action, w, env)) * x

            else:
                next_action = epsilon_greedy_action_w(
                    env, w, next_state, epsilon(episode), rng=rng
                )
                w += alpha(episode) * (reward + gamma * q_hat(next_state, next_action, w, env) - q_hat(state, action, w, env)) * x

            if verbose:
                print (f"Episode {episode+1}, Step {i+1}: S={state}, A={action}, R={reward}, S'={next_state}, w={w}")

            state = next_state
            action = next_action
            i += 1

            if history:
                G += reward * (gamma ** (i-1))

        if history:
            ws.append(w.copy())
            returns.append(G)
            ep_lens.append(i)

    if history:
        return w, {'ws': ws, 'returns': returns, 'ep_lens': ep_lens}

    return w  

create_fourier_basis_coefs

create_fourier_basis_coefs(dim, order)

Create Fourier basis coefficients for given dimension and order. param dim: dimension of the state features param order: order of the Fourier basis

Source code in gym_classics2/algorithms/linear_approximation.py
def create_fourier_basis_coefs(dim, order): 
    """ Create Fourier basis coefficients for given dimension and order. 
        param dim: dimension of the state features
        param order: order of the Fourier basis
    """  
    return np.array(list(product(range(order+1), repeat=dim)))

transformation_fourier_basis

transformation_fourier_basis(min, max, order)

Create a Fourier basis transformation function for given min/max ranges and order.

To use this transformation with semi_gradient_Sarsa you need to overwrite the state_features function like this:

def state_features(s, env): return trans_fb(env.decode(s)) gym_classics2.algorithms.linear_approximation.state_features = state_features

param min: minimum values for each dimension param max: maximum values for each dimension param order: order of the Fourier basis

Source code in gym_classics2/algorithms/linear_approximation.py
def transformation_fourier_basis(min, max, order):
    """ Create a Fourier basis transformation function for given min/max ranges and order.

        To use this transformation with semi_gradient_Sarsa you need to overwrite the state_features 
        function like this:

        def state_features(s, env): return trans_fb(env.decode(s))
        gym_classics2.algorithms.linear_approximation.state_features = state_features

        param min: minimum values for each dimension
        param max: maximum values for each dimension
        param order: order of the Fourier basis
    """  

    min = np.array(min)
    max = np.array(max)
    coefs = create_fourier_basis_coefs(len(min), order)

    def fourier_basis(s):
        # normalize state to [0,1]
        s = np.array(s)
        assert s.shape == min.shape, "State dimension does not match Fourier basis dimension"

        norm_s = (s - min) / (max - min)
        return np.cos(np.pi * np.dot(coefs, norm_s))

    return fourier_basis

Eligibility traces

gym_classics2.algorithms.eligibility_traces

This file implements the semi-gradient SARSA(lambda) algorithm for control with linear function approximation and eligibility traces. The user needs to implement the state_features function to convert states to feature vectors.

active_weights

active_weights(a, sf_len)

helper for q_hat()

Source code in gym_classics2/algorithms/eligibility_traces.py
def active_weights(a, sf_len):
    """helper for q_hat()"""
    return [0] + list(range(a*sf_len+1, a*sf_len+sf_len+1))

semi_gradient_Sarsa_lambda

semi_gradient_Sarsa_lambda(env, n, epsilon, alpha, gamma, lam, w=None, max_episode_length=1000, verbose=False, history=False, rng=None)

Semi-gradient SARSA(lambda): on-policy control with linear function approximation and eligibility traces.

Parameters

env : GymClassicsBaseEnv Episodic environment used to generate experience. n : int Number of episodes. epsilon : float Exploration rate for epsilon-greedy policy. alpha : float Step size. gamma : float Discount factor. lam : float Trace-decay parameter lambda in [0, 1]. w : array-like or None Initial weights. If None, initializes to zeros. max_episode_length : int Maximum number of steps per episode. verbose : bool Whether to print step-by-step diagnostics. rng : numpy.random.Generator or int or None Random generator or seed for exploration and tie-breaking.

Returns

w : np.ndarray Learned weight vector.

Source code in gym_classics2/algorithms/eligibility_traces.py
def semi_gradient_Sarsa_lambda(
    env,
    n,
    epsilon,
    alpha,
    gamma,
    lam,
    w=None,
    max_episode_length=1000,
    verbose=False,
    history=False,
    rng=None,
):
    """
    Semi-gradient SARSA(lambda): on-policy control with linear function approximation
    and eligibility traces.

    Parameters
    ----------
    env : GymClassicsBaseEnv
        Episodic environment used to generate experience.
    n : int
        Number of episodes.
    epsilon : float
        Exploration rate for epsilon-greedy policy.
    alpha : float
        Step size.
    gamma : float
        Discount factor.
    lam : float
        Trace-decay parameter lambda in [0, 1].
    w : array-like or None
        Initial weights. If None, initializes to zeros.
    max_episode_length : int
        Maximum number of steps per episode.
    verbose : bool
        Whether to print step-by-step diagnostics.
    rng : numpy.random.Generator or int or None
        Random generator or seed for exploration and tie-breaking.

    Returns
    -------
    w : np.ndarray
        Learned weight vector.
    """

    assert gamma >= 0 and gamma <= 1, "gamma must be in [0,1]"
    assert lam >= 0 and lam <= 1, "lambda must be in [0,1]"
    assert n > 0, "number of episodes must be positive"
    assert max_episode_length > 0, "max episode length must be positive"

    rng = get_rng(rng)

    if not isinstance(alpha, Schedule):
        alpha = ConstantSchedule(alpha)
    if not isinstance(epsilon, Schedule):
        epsilon = ConstantSchedule(epsilon)

    if w is None:
        state, _ = env.reset()
        w = np.zeros(len(state_action_features(state, 0, env)))

    if history:
        ws = []
        ws.append(w.copy())
        returns = []
        ep_lens = []


    for episode in tqdm(range(n), desc="Semi-Gradient SARSA(lambda)", disable=verbose):
        state, _ = env.reset()
        action = epsilon_greedy_action_w(
            env, w, state, epsilon(episode), rng=rng
        )

        # eligibility trace vector, same size as w
        z = np.zeros_like(w)
        Q_old = 0

        done = False
        i = 0

        G = 0  # for tracking returns if history is enabled

        while not done and i < max_episode_length:
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            G += reward * (gamma ** i)  # accumulate return if history is enabled

            # current feature vector for (state, action)
            x = state_action_features(state, action, env)

            # update trace
            z = gamma * lam * z + (1 - alpha(episode) * gamma * lam * np.dot(z, x)) * x

            if terminated:
                delta = reward - q_hat(state, action, w, env)
            else:
                next_action = epsilon_greedy_action_w(
                    env, w, next_state, epsilon(episode), rng=rng
                )
                delta = reward + gamma * q_hat(next_state, next_action, w, env) - q_hat(state, action, w, env)

            # semi-gradient weight update
            Q = q_hat(state, action, w, env)
            Q_prime = q_hat(next_state, next_action, w, env) if not terminated else 0            
            w += alpha(episode) * (delta + Q - Q_old) * z - alpha(episode) * (Q - Q_old) * x

            Q_old = Q_prime

            if verbose:
                if terminated:
                    print(
                        f"Episode {episode+1}, Step {i+1}: "
                        f"S={state}, A={action}, R={reward}, S'={next_state}, "
                        f"delta={delta}, z={z}, w={w}"
                    )
                else:
                    print(
                        f"Episode {episode+1}, Step {i+1}: "
                        f"S={state}, A={action}, R={reward}, S'={next_state}, A'={next_action}, "
                        f"delta={delta}, z={z}, w={w}"
                    )

            if done:
                break

            state = next_state
            action = next_action
            i += 1

        if history:
            returns.append(G)
            ws.append(w.copy())
            ep_lens.append(i)


    if history:        
        return w, {'ws': ws, 'returns': returns, 'ep_lens': ep_lens}

    return w

Policy-gradient methods

gym_classics2.algorithms.policy_gradient_methods

This file implements policy gradient methods for learning parameterized policies. The main algorithm implemented is REINFORCE, which is a Monte Carlo policy gradient method that updates policy parameters based on the returns observed in sampled episodes. The policy is represented using a softmax function over linear state-action features, and the algorithm estimates the policy gradient using the log-likelihood of actions taken in the episodes. This implementation allows for learning stochastic policies that can handle exploration and exploitation in reinforcement learning tasks.

The user has to overwrite the state_features function to convert state ids into feature vectors suitable for the environment being used.

h

h(s, a, theta, env)

Return the linear action preference for state s and action a.

Source code in gym_classics2/algorithms/policy_gradient_methods.py
def h(s,a,theta,env):
    """Return the linear action preference for state ``s`` and action ``a``."""
    return np.dot(theta, state_action_features(s,a,env))

pi

pi(s, theta, env)

Return the softmax action-probability vector for a state.

Source code in gym_classics2/algorithms/policy_gradient_methods.py
def pi(s,theta,env):
    """Return the softmax action-probability vector for a state."""
    hs = np.array([h(s,a,theta,env) for a in range(env.action_space.n)])
    exp_hs = np.exp(hs)
    return exp_hs / np.sum(exp_hs)

sample_episode_approx_policy

sample_episode_approx_policy(env, pi, theta, max_episode_length=1000, rng=None)

Sample an episode using the policy defined by pi and theta.

:param env: the environment to sample from :param pi: the policy function that takes state, theta, and env as input and returns a probability distribution over actions :param theta: the policy parameters :param max_episode_length: maximum number of steps to sample in the episode :param rng: NumPy generator or integer seed used to sample actions :return: a list of (state, action, reward, next_state) tuples representing the sampled episode

Source code in gym_classics2/algorithms/policy_gradient_methods.py
def sample_episode_approx_policy(env, pi, theta, max_episode_length=1000, rng=None):
    """
    Sample an episode using the policy defined by pi and theta.

    :param env: the environment to sample from
    :param pi: the policy function that takes state, theta, and env as input and
        returns a probability distribution over actions
    :param theta: the policy parameters
    :param max_episode_length: maximum number of steps to sample in the episode
    :param rng: NumPy generator or integer seed used to sample actions
    :return: a list of (state, action, reward, next_state) tuples representing the sampled episode
    """

    rng = get_rng(rng)
    s, _ = env.reset()
    episode_data = []

    for t in range(max_episode_length):
        a = rng.choice(env.action_space.n, p=pi(s, theta, env))
        next_s, r, done, _, _ = env.step(a)
        episode_data.append((s, a, r, next_s))
        s = next_s

        if done:
            break

    return episode_data

choose_action_w

choose_action_w(env, pi, theta, state, rng=None)

Choose an action based on the policy defined by pi and theta for the given state.

:param env: the environment :param pi: the policy function that takes state, theta, and env as input and returns a probability distribution over actions :param theta: the policy parameters :param state: the current state :param rng: NumPy generator or integer seed used to sample the action :return: the chosen action

Source code in gym_classics2/algorithms/policy_gradient_methods.py
def choose_action_w(env, pi, theta, state, rng=None):
    """
    Choose an action based on the policy defined by pi and theta for the given state.

    :param env: the environment
    :param pi: the policy function that takes state, theta, and env as input and
        returns a probability distribution over actions 
    :param theta: the policy parameters
    :param state: the current state
    :param rng: NumPy generator or integer seed used to sample the action
    :return: the chosen action
    """
    rng = get_rng(rng)
    return rng.choice(env.action_space.n, p=pi(state, theta, env))

REINFORCE

REINFORCE(env, n, alpha, gamma, theta=None, max_episode_length=1000, verbose=False, history=False, rng=None)

REINFORCE: Monte Carlo policy gradient method with linear function approximation. Parameters


env : GymClassicsBaseEnv Episodic environment used to generate experience. n : int Number of episodes. alpha : float Step size. gamma : float Discount factor. theta : array-like or None Initial policy parameters. If None, initializes to zeros. max_episode_length : int Maximum number of steps per episode. verbose : bool Whether to print step-by-step diagnostics. history : bool Whether to return learning history (returns, episode lengths, parameter values).
rng : numpy.random.Generator or int or None Random generator or seed used to sample actions.

Source code in gym_classics2/algorithms/policy_gradient_methods.py
def REINFORCE(
    env,
    n,
    alpha,
    gamma,
    theta = None,
    max_episode_length=1000,
    verbose=False,
    history=False,
    rng=None,
    ):
    """REINFORCE: Monte Carlo policy gradient method with linear function approximation.
    Parameters
    ----------
    env : GymClassicsBaseEnv
        Episodic environment used to generate experience.
    n : int
        Number of episodes.
    alpha : float
        Step size.
    gamma : float
        Discount factor.
    theta : array-like or None
        Initial policy parameters. If None, initializes to zeros.
    max_episode_length : int
        Maximum number of steps per episode.
    verbose : bool
        Whether to print step-by-step diagnostics.
    history : bool
        Whether to return learning history (returns, episode lengths, parameter values).    
    rng : numpy.random.Generator or int or None
        Random generator or seed used to sample actions.
    """

    assert gamma >= 0 and gamma <= 1, "gamma must be in [0  ,1]"
    assert n > 0, "number of episodes must be positive"
    assert max_episode_length > 0, "max episode length must be positive"

    rng = get_rng(rng)

    if isinstance(env.observation_space, gym.spaces.Discrete):
        warnings.warn("The environment has a discrete state space. Consider using a tabular method instead of function approximation.")

    if isinstance(alpha, float):
        alpha = ConstantSchedule(alpha)

    if theta is None:
        state, _ = env.reset()
        theta = np.zeros(len(state_action_features(state, 0, env)))

    if history:
        returns = []        
        ep_lens = []
        thetas = []
        thetas.append(theta.copy())

    for episode in tqdm(range(n), desc="Episodes", disable=verbose):
        if verbose:
            print(f"Episode {episode+1}/{n}")

        # sample complete episode using pi (this is a MC method)
        episode_data = sample_episode_approx_policy(
            env, pi, theta, max_episode_length, rng=rng
        )

        for t in range(len(episode_data)):
            # update policy for each step in the episode using the return observed from that step on
            #print(episode_data[t])

            G = np.sum([e[2] for e in episode_data[t:]] * (gamma ** np.arange(len(episode_data[t:]))))
            if history and t == 0:
                returns.append(G)   

            s,a,r,next_s = episode_data[t]

            # ln policy gradient= x(s,a)- sum_b pi(b|s,theta) x(s,b)
            grad_log_pi = state_action_features(s, a, env) - sum([pi(s, theta, env)[b] * state_action_features(s, b, env) for b in range(env.action_space.n)])

            if verbose: 
                print (f"t: {t}, G: {G:.2f}, grad_log_pi: {grad_log_pi}")

            theta += alpha(episode) * (gamma**t) * G * grad_log_pi

        if history:
            thetas.append(theta.copy())
            ep_lens.append(len(episode_data))

    if history:
        return theta, {'thetas': thetas, 'returns': returns, 'ep_lens': ep_lens}

    return theta

AC

AC(env, n, alpha_policy, alpha_value, gamma, max_episode_length=1000, verbose=False, history=False, rng=None)

Actor-Critic: Policy gradient method with linear function approximation and TD learning for the value function.

Parameters

env : Episodic environment used to generate experience. n : int Number of episodes. alpha_policy : float or Schedule Step size for policy updates. If a float is provided, it will be converted to a ConstantSchedule. If a Schedule is provided, it will be used directly. alpha_value : float or Schedule Step size for value function updates. If a float is provided, it will be converted to a ConstantSchedule. If a Schedule is provided, it will be used directly. gamma : float Discount factor. max_episode_length : int Maximum number of steps per episode. verbose : bool Whether to print step-by-step diagnostics. history : bool Whether to return learning history (returns, episode lengths, parameter values). rng : numpy.random.Generator or int or None Random generator or seed used to sample actions.

Returns
-------
theta : array-like
    Final policy parameters after training.
w : array-like
    Final value function parameters after training.
history : dict (optional)
    If history=True, a dictionary containing the learning history with keys: 
        'thetas': list of policy parameter vectors at each episode,
        'ws': list of value function parameter vectors at each episode,
        'returns': list of returns observed at the end of each episode,
        'ep_lens': list of episode lengths (number of steps) for each episode.
Source code in gym_classics2/algorithms/policy_gradient_methods.py
def AC(
    env,
    n,
    alpha_policy,
    alpha_value,
    gamma,
    max_episode_length=1000,
    verbose=False,
    history=False,
    rng=None,
    ):
    """Actor-Critic: Policy gradient method with linear function approximation and TD learning for the value function.

    Parameters
    ----------
    env : Episodic environment used to generate experience.
    n : int
        Number of episodes.
    alpha_policy : float or Schedule
        Step size for policy updates. If a float is provided, it will be converted to a ConstantSchedule. If a Schedule is provided, it will be used directly.
    alpha_value : float or Schedule
        Step size for value function updates. If a float is provided, it will be converted to a ConstantSchedule. If a Schedule is provided, it will be used directly. 
    gamma : float
        Discount factor.
    max_episode_length : int
        Maximum number of steps per episode.
    verbose : bool
        Whether to print step-by-step diagnostics.
    history : bool
        Whether to return learning history (returns, episode lengths, parameter values).
    rng : numpy.random.Generator or int or None
        Random generator or seed used to sample actions.

        Returns
        -------
        theta : array-like
            Final policy parameters after training.
        w : array-like
            Final value function parameters after training.
        history : dict (optional)
            If history=True, a dictionary containing the learning history with keys: 
                'thetas': list of policy parameter vectors at each episode,
                'ws': list of value function parameter vectors at each episode,
                'returns': list of returns observed at the end of each episode,
                'ep_lens': list of episode lengths (number of steps) for each episode.
        """

    assert gamma >= 0 and gamma <= 1, "gamma must be in [0  ,1]"
    assert n > 0, "number of episodes must be positive"
    assert max_episode_length > 0, "max episode length must be positive"

    rng = get_rng(rng)

    if isinstance(env.observation_space, gym.spaces.Discrete):
        warnings.warn("The environment has a discrete state space. Consider using a tabular method instead of function approximation.")

    if isinstance(alpha_policy, float):
        alpha_policy = ConstantSchedule(alpha_policy)
    if isinstance(alpha_value, float):
        alpha_value = ConstantSchedule(alpha_value) 

    state, _ = env.reset()

    # for simplicity we use the same features for the value function and the policy approximation
    # value function weights
    w = np.zeros(len(state_features(state, env)))
    # policy weights
    theta = np.zeros(len(state_action_features(state, 0, env)))

    if history:
        returns = []        
        ep_lens = []
        ws = []
        ws.append(w.copy())
        thetas = []
        thetas.append(theta.copy())

    for episode in tqdm(range(n), desc="Episodes", disable=verbose):
        if verbose:
            print(f"Episode {episode+1}/{n}")

        disc_factor = 1.0
        state, _ = env.reset()
        done = False
        i = 0
        G = 0.0

        while not done and  i < max_episode_length:
            # use actor to determine next action
            a = rng.choice(env.action_space.n, p=pi(state, theta, env))

            # execute action
            next_state, r, done, _, _ = env.step(a)

            # use critic to calculate  TD error
            td_error = r + gamma * np.dot(w, state_features(next_state, env)) - np.dot(w, state_features(state, env))

            # update critic
            w += alpha_value(episode) * td_error * state_features(state, env)

            # update actor
            grad_log_pi = state_action_features(state, a, env) - sum([pi(state, theta, env)[b] * state_action_features(state, b, env) for b in range(env.action_space.n)])
            theta += alpha_policy(episode) * disc_factor * td_error * grad_log_pi

            G += disc_factor * r
            disc_factor *= gamma      
            state = next_state
            i += 1

        if history:
            thetas.append(theta.copy())
            ws.append(w.copy())
            returns.append(G)
            ep_lens.append(i)

    if history:
        return theta, w, {'thetas': thetas, 'returns': returns, 'ep_lens': ep_lens}

    return theta, w

Schedules

gym_classics2.algorithms.schedules

This file implements different schedules to reduce the learning rate alpha or the exploration rate epsilon over time (used in on-policy methods for GLIE). These are commonly used in reinforcement learning algorithms to improve convergence and performance.

Schedule

Base class for schedules.

Source code in gym_classics2/algorithms/schedules.py
class Schedule:
    """Base class for schedules."""
    def __call__(self, t):
        """Returns the scheduled value at time step (or episode) t."""
        raise NotImplementedError

__call__

__call__(t)

Returns the scheduled value at time step (or episode) t.

Source code in gym_classics2/algorithms/schedules.py
def __call__(self, t):
    """Returns the scheduled value at time step (or episode) t."""
    raise NotImplementedError

ConstantSchedule

Bases: Schedule

A schedule that always returns a constant value.

Source code in gym_classics2/algorithms/schedules.py
class ConstantSchedule(Schedule):
    """A schedule that always returns a constant value."""
    def __init__(self, value):
        self.value = float(value)

    def __call__(self, t):
        return self.value

StepSchedule

Bases: Schedule

A schedule that always returns a constant value.

Source code in gym_classics2/algorithms/schedules.py
class StepSchedule(Schedule):
    """A schedule that always returns a constant value."""
    def __init__(self, high_value, low_value, steps):
        self.high_value = float(high_value)
        self.low_value = float(low_value)
        self.steps = int(steps)

    def __call__(self, t):
        if t < self.steps:
            return self.high_value
        return self.low_value

LinearDecaySchedule

Bases: Schedule

A schedule that decays linearly from initial_value to min_value over decay_steps.

Source code in gym_classics2/algorithms/schedules.py
class LinearDecaySchedule(Schedule):
    """A schedule that decays linearly from initial_value to min_value over decay_steps."""
    def __init__(self, initial_value, min_value, decay_steps):
        self.initial_value = float(initial_value)
        self.min_value = float(min_value)
        self.decay_steps = int(decay_steps)

    def __call__(self, t):
        fraction = min(float(t) / max(1, self.decay_steps), 1.0)
        return self.initial_value + fraction * (self.min_value - self.initial_value)

ExponentialDecaySchedule

Bases: Schedule

A schedule that decays exponentially with a given decay_rate. Value at time t is max(min_value, initial_value * (decay_rate ** t)).

Source code in gym_classics2/algorithms/schedules.py
class ExponentialDecaySchedule(Schedule):
    """A schedule that decays exponentially with a given decay_rate.
    Value at time t is max(min_value, initial_value * (decay_rate ** t)).
    """
    def __init__(self, initial_value, min_value, decay_rate):
        self.initial_value = float(initial_value)
        self.min_value = float(min_value)
        self.decay_rate = float(decay_rate)

    def __call__(self, t):
        return max(self.min_value, self.initial_value * (self.decay_rate ** t))

InverseDecaySchedule

Bases: Schedule

A schedule that decays inversely proportional to t. Value at time t is max(min_value, initial_value / t) for t > 0.

Source code in gym_classics2/algorithms/schedules.py
class InverseDecaySchedule(Schedule):
    """A schedule that decays inversely proportional to t.
    Value at time t is max(min_value, initial_value / t) for t > 0.
    """
    def __init__(self, initial_value, min_value=0.0):
        self.initial_value = float(initial_value)
        self.min_value = float(min_value)

    def __call__(self, t):
        if t == 0:
            return self.initial_value
        return max(self.min_value, self.initial_value / t)

plot_schedule

plot_schedule(schedule, steps=1000)

Plots the values of a schedule over a given number of steps.

Example

from gym_classics2.algorithms.schedules import ConstantSchedule, LinearDecaySchedule, ExponentialDecaySchedule, InverseDecaySchedule, plot_schedule

Create different schedules

constant_sched = ConstantSchedule(0.1) linear_sched = LinearDecaySchedule(initial_value=1.0, min_value=0.1, decay_steps=1000) exponential_sched = ExponentialDecaySchedule(initial_value=1.0, min_value=0.1, decay_rate=0.99) inverse_sched = InverseDecaySchedule(initial_value=10.0, min_value=0.1)

Plot them

plot_schedule(constant_sched, steps=500) plot_schedule(linear_sched, steps=1000) plot_schedule(exponential_sched, steps=1000) plot_schedule(inverse_sched, steps=1000)

Parameters:

Name Type Description Default
schedule

A Schedule instance (or any callable taking an integer step).

required
steps

The number of steps to plot.

1000
Source code in gym_classics2/algorithms/schedules.py
def plot_schedule(schedule, steps=1000):
    """Plots the values of a schedule over a given number of steps.

    Example:
        from gym_classics2.algorithms.schedules import ConstantSchedule, LinearDecaySchedule, ExponentialDecaySchedule, InverseDecaySchedule, plot_schedule

        # Create different schedules
        constant_sched = ConstantSchedule(0.1)
        linear_sched = LinearDecaySchedule(initial_value=1.0, min_value=0.1, decay_steps=1000)
        exponential_sched = ExponentialDecaySchedule(initial_value=1.0, min_value=0.1, decay_rate=0.99)
        inverse_sched = InverseDecaySchedule(initial_value=10.0, min_value=0.1)

        # Plot them
        plot_schedule(constant_sched, steps=500)
        plot_schedule(linear_sched, steps=1000)
        plot_schedule(exponential_sched, steps=1000)
        plot_schedule(inverse_sched, steps=1000)

    Args:
        schedule: A Schedule instance (or any callable taking an integer step).
        steps: The number of steps to plot.
    """
    import matplotlib.pyplot as plt

    values = [schedule(t) for t in range(steps)]

    plt.figure(figsize=(8, 5))
    plt.plot(range(steps), values, linewidth=2)
    plt.xlabel('Step')
    plt.ylabel('Schedule Value')
    plt.title('Schedule Plot')
    plt.grid(True)
    plt.show()