cartpole.py

"""Cartpole controller that interacts with Blender environment through OpenAI.

Run
    python cartpole.py
from this directory.

Note, the custom environment is registered in `cartpole_gym/__init__.py`. The
file `cartpole_gym/cartpole_env.py` defines a OpenAI compatible Blender environment.
The actual environment logic is implemented in Blender script
`cartpole_gym/cartpole.blend.py` and scene `cartpole_gym/cartpole.blend`.
"""


import gym
import cartpole_gym

KAPPA = 30


def control(obs):
    # Simple P controller defining the error as xpole-xcart
    xcart, xpole, _ = obs
    return (xpole - xcart) * KAPPA


def main():
    # Create the environment. The environment is registered in
    # `cartpole_gym/__init__.py`. Set `real_time=True` when you
    # want to keep the simulation running until the agent response.
    env = gym.make("blendtorch-cartpole-v0", real_time=False)

    obs = env.reset()
    while True:
        obs, reward, done, info = env.step(control(obs))
        env.render()
        if done:
            obs = env.reset()
    env.done()


if __name__ == "__main__":
    main()