Train the inverted double link cartpole model

We will now demonstrate training the model using StableBaselines 3

pip install stable-baselines3
pip install gymnasium

import os
from env import TwoLinkCartPoleEnv
import gymnasium as gym
from stable_baselines3 import SAC
from stable_baselines3.common.vec_env import DummyVecEnv

gym.register(
    id="kdFlex-2link-CartPole",
    entry_point=TwoLinkCartPoleEnv,
    max_episode_steps=1000,
)

# use render_mode="human" if you want to visualize training
train_env = gym.make("kdFlex-2link-CartPole", render_mode="none")
train_env = DummyVecEnv([lambda: train_env])

model = SAC(
    policy="MlpPolicy",
    env=train_env,
)

# check if this is running in a test environment
if os.getenv("DTEST_RUNNING", True):
    # in a testing environment, run for less
    model.learn(total_timesteps=1000)
else:
    # this will take around 50 minutes
    model.learn(total_timesteps=400000)

train_env.close()
test_env = gym.make("kdFlex-2link-CartPole", render_mode="human", sync_real_time=True)
test_env = DummyVecEnv([lambda: test_env])

[WebUI] Listening at http://newton:34393

terminated = False
obs = test_env.reset()
while not terminated:
    action, _states = model.predict(obs)
    obs, reward, terminated, info = test_env.step(action)

model = SAC.load("./models/demo_model", env=test_env)

# cleanup
test_env.close()

Summary

Amazing! You can now implement a custom gymnasium environment using kdFlex. By solving

Further Readings