pykesko 0.0.4

Robotics Simulator written in Rust

PyKesko

PyKesko is the python package for the Kesko simulator. It contains python bindings to the Rust code and provides an API to interact with simulator.

Note: PyKesko is currently pre-release, unstable and very limited.

Getting Started

Prerequisites

Make sure you have one of the supported python versions installed (3.7-3.11). It is also recommended to use a virtual environment.

Installation

pip install pykesko

Usage

Run the simulator

Run and step the simulator

from pykesko import Kesko


if __name__ == "__main__":
    kesko = Kesko()
    kesko.initialize()

    for _ in range(1000):
        kesko.step()

    kesko.close()

Spawn bodies

Kesko has som built in models that can be spawned.

from pykesko import Kesko, KeskoModel
from pykesko.protocol.commands import Spawn
from pykesko.color import Color


if __name__ == "__main__":
    kesko = Kesko()
    kesko.initialize()

    kesko.send(
        [
            Spawn(model=KeskoModel.Plane, position=[0.0, 0.0, 0.0], color=Color.WHITE),
            Spawn(model=KeskoModel.Humanoid, position=[0.0, 2.0, 0.0], color=Color.DEEP_ORANGE),
        ]
    )

    for _ in range(1000):
        kesko.step()

    kesko.close()

Use a Gym environment

In the future PyKesko will contain a bunch of gym environments. Currently there is only one prototype environment consisting of a four legged agent which is tasked with learning to walk.

import gym

env = gym.make("pykesko:Spider-v0", render_mode="human")
env.reset()

for _ in range(1000):
    env.step(env.action_space.sample())

env.close()

Use Stable Baselines3 to train an agent to walk

Since PyKesko has Gym integration it is really easy to train with a library like Stable Baselines3. Make sure to use the bleeding edge version of Stable Baselines3 that has support for the latest Gym version

Install Stable Baselines3 bleeding edge version

pip install git+https://github.com/carlosluis/stable-baselines3@fix_tests

Train agent using PPO

import gym
from stable_baselines3 import PPO


env = gym.make("pykesko:Spider-v0", render_mode="human")

model = PPO("MlpPolicy", env)
model.learn(total_timesteps=100_000, progress_bar=True)

env.close()

Running PPO for about 5 million steps yield the result below. Note that in order to train in reasonable amount of time you should use parallel non-render environments.

Train parallel environments in headless mode

Below is an example how to use multiple parallel environments without rendering (headless mode). It also includes saving checkpoints as well as logging to tensorboard

from datetime import datetime

import gym

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import CheckpointCallback
from stable_baselines3.common.vec_env import SubprocVecEnv
from stable_baselines3.common.utils import set_random_seed


ENV = "Spider-v0"
FULL_ENV = f"pykesko:Spider-v0"
MODEL_NAME = f"{ENV}_PPO"

# Number of environments, should be close to the number of cpu's available
NUM_ENVS = 16

# PPO parameters
N_STEPS = 2048 // NUM_ENVS  # Number of steps to take for each PPO iteration
N_EPOCHS = 10  # Number of epochs to use when optimizing the surrogate loss
BATCH_SIZE = 64
TOTAL_TIME_STEPS = 8_000_000  # Total number of time steps to train for

SAVE_EVERY = 200_000  # How often to save a checkpoint of the model
SAVE_FOLDER = f"./runs/{MODEL_NAME}_{datetime.now()}"


def make_env(env_id: str, rank: int, seed: int = 0):
    """
    Utility function for multiprocessed env.

    :param env_id: the environment ID
    :param seed: the initial seed for RNG
    :param rank: index of the subprocess
    """

    def _init():
        env = gym.make(env_id)
        env.seed(seed + rank)
        return env

    set_random_seed(seed)
    return _init


if __name__ == "__main__":

    # Create NUM_ENVS environments
    env = SubprocVecEnv([make_env(env_id=FULL_ENV, rank=i) for i in range(NUM_ENVS)])

    checkpoint_callback = CheckpointCallback(
        save_freq=max(SAVE_EVERY // NUM_ENVS, 1),
        save_path=f"{SAVE_FOLDER}/model_checkpoints",
        name_prefix=MODEL_NAME,
    )

    model = PPO(
        "MlpPolicy",
        env,
        tensorboard_log=f"{SAVE_FOLDER}/tb_logs",
        n_steps=N_STEPS,
        n_epochs=N_EPOCHS,
        batch_size=BATCH_SIZE,
    )

    model.learn(
        total_timesteps=TOTAL_TIME_STEPS,
        callback=checkpoint_callback,
        progress_bar=True,
    )

    env.close()