joystick_control.py

#!/usr/bin/env python3

"""
This script allows you to manually control the simulator or Duckiebot
using a Logitech Game Controller, as well as record trajectories.
"""

import argparse
import json
import sys

import gym
import numpy as np
import pyglet
from pyglet.window import key

from gym_duckietown.envs import DuckietownEnv

parser = argparse.ArgumentParser()
parser.add_argument("--env-name", default=None)
parser.add_argument("--map-name", default="udem1")
parser.add_argument("--distortion", default=False, action="store_true")
parser.add_argument("--draw-curve", action="store_true", help="draw the lane following curve")
parser.add_argument("--domain-rand", action="store_true", help="enable domain randomization")
args = parser.parse_args()

if args.env_name is None:
    env = DuckietownEnv(
        map_name=args.map_name, distortion=args.distortion, domain_rand=args.domain_rand, max_steps=np.inf
    )
else:
    env = gym.make(args.env_name)

env.reset()
env.render()

# global variables for demo recording
positions = []
actions = []
demos = []
recording = False


def write_to_file(demos):
    num_steps = 0
    for demo in demos:
        num_steps += len(demo["actions"])
    print("num demos:", len(demos))
    print("num steps:", num_steps)

    # Store the trajectories in a JSON file
    with open("experiments/demos_{}.json".format(args.map_name), "w") as outfile:
        json.dump({"demos": demos}, outfile)


def process_recording():
    global positions, actions, demos

    if len(positions) == 0:
        # Nothing to delete
        if len(demos) == 0:
            return

        # Remove the last recorded demo
        demos.pop()
        write_to_file(demos)
        return

    p = list(map(lambda p: [p[0].tolist(), p[1]], positions))
    a = list(map(lambda a: a.tolist(), actions))

    demo = {"positions": p, "actions": a}

    demos.append(demo)

    # Write all demos to this moment
    write_to_file(demos)


@env.unwrapped.window.event
def on_key_press(symbol, modifiers):
    """
    This handler processes keyboard commands that
    control the simulation
    """

    if symbol == key.BACKSPACE or symbol == key.SLASH:
        print("RESET")
        env.reset()
        env.render()
    elif symbol == key.PAGEUP:
        env.unwrapped.cam_angle[0] = 0
        env.render()
    elif symbol == key.ESCAPE:
        env.close()
        sys.exit(0)


@env.unwrapped.window.event
def on_joybutton_press(joystick, button):
    """
    Event Handler for Controller Button Inputs
    Relevant Button Definitions:
    1 - A - Starts / Stops Recording
    0 - X - Deletes last Recording
    2 - Y - Resets Env.

    Triggers on button presses to control recording capabilities
    """
    global recording, positions, actions

    # A Button
    if button == 1:
        if not recording:
            print("Start recording, Press A again to finish")
            recording = True
        else:
            recording = False
            process_recording()
            positions = []
            actions = []
            print("Saved recording")

    # X Button
    elif button == 0:
        recording = False
        positions = []
        actions = []
        process_recording()
        print("Deleted last recording")

    # Y Button
    elif button == 3:
        print("RESET")
        env.reset()
        env.render()

    # Any other button thats not boost prints help
    elif button != 5:
        helpstr1 = "A - Starts / Stops Recording\nX - Deletes last Recording\n"
        helpstr2 = "Y - Resets Env.\nRB - Hold for Boost"

        print("Help:\n{}{}".format(helpstr1, helpstr2))


def update(dt):
    """
    This function is called at every frame to handle
    movement/stepping and redrawing
    """
    global recording, positions, actions

    # No actions took place
    if round(joystick.x, 2) == 0.0 and round(joystick.y, 2) == 0.0:
        return

    x = round(joystick.y, 2)
    z = round(joystick.x, 2)

    action = np.array([-x, -z])

    # Right trigger, speed boost
    if joystick.buttons[5]:
        action *= 1.5

    if recording:
        positions.append((env.unwrapped.cur_pos, env.unwrapped.cur_angle))
        actions.append(action)

    obs, reward, done, info = env.step(action)
    print("step_count = %s, reward=%.3f" % (env.unwrapped.step_count, reward))

    if done:
        print("done!")
        env.reset()
        env.render()

        if recording:
            process_recording()
            positions = []
            actions = []
            print("Saved Recoding")

    env.render()


pyglet.clock.schedule_interval(update, 1.0 / env.unwrapped.frame_rate)

# Registers joysticks and recording controls
joysticks = pyglet.input.get_joysticks()
assert joysticks, "No joystick device is connected"
joystick = joysticks[0]
joystick.open()
joystick.push_handlers(on_joybutton_press)

# Enter main event loop
pyglet.app.run()

env.close()