mirror of
https://github.com/wassname/pyrobolearn.git
synced 2026-09-09 11:31:38 +08:00
88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
"""Imitation learning demonstration using DMPs and ROS with the Franka robot.
|
|
"""
|
|
|
|
import pyrobolearn as prl
|
|
|
|
|
|
# variables
|
|
joint_ids = None # None for all the actuated joints, or you can select which joint you want to move; e.g. [0, 1, 2]
|
|
num_basis = 20
|
|
rate = 30
|
|
use_real_robot = True
|
|
|
|
# create middleware
|
|
ros = prl.middlewares.ROS(subscribe=True, teleoperate=True)
|
|
if not use_real_robot:
|
|
ros.disable_middleware() # disable the middleware if not using the real robot
|
|
|
|
# create simulator
|
|
sim = prl.simulators.Bullet(middleware=ros)
|
|
|
|
# create basic world (with gravity and floor)
|
|
world = prl.worlds.BasicWorld(sim)
|
|
|
|
# load Franka Emika Panda robot in the world
|
|
input("Press Enter to load the robot")
|
|
robot = prl.robots.Franka(sim)
|
|
world.load_robot(robot)
|
|
robot.print_info()
|
|
|
|
# create state/action
|
|
state = prl.states.ExponentialPhaseState(ticks=rate)
|
|
action = prl.actions.JointPositionAction(robot, joint_ids=joint_ids)
|
|
print("State: {}".format(state))
|
|
print("Action: {}".format(action))
|
|
|
|
# create environment
|
|
env = prl.envs.Env(world, state)
|
|
|
|
# create DMP policy
|
|
policy = prl.policies.BioDiscreteDMPPolicy(action, state, num_basis=num_basis, rate=rate)
|
|
|
|
# create interface/bridge
|
|
interface = prl.interfaces.MouseKeyboardInterface(sim)
|
|
bridge = prl.bridges.BridgeMouseKeyboardImitationTask(world, interface=interface, verbose=True)
|
|
|
|
# create recorder
|
|
recorder = prl.recorders.StateRecorder(prl.states.JointPositionState(robot, joint_ids=joint_ids), rate=rate)
|
|
|
|
# create imitation learning task
|
|
task = prl.tasks.ILTask(env, policy, interface=bridge, recorders=recorder)
|
|
|
|
# record, train, and test policy using the policy
|
|
# task.run()
|
|
|
|
# record demonstrations in simulation/reality
|
|
print("\nRecording phase: press `ctrl+r` to start/stop the recording. Once finished, press `shift+r`.")
|
|
task.record(signal_from_interface=True)
|
|
print("Recording phase: finished the recording!")
|
|
sim.disable_middleware() # disable the middleware (get/set info only from/to simulation)
|
|
|
|
# train policy
|
|
input("Press Enter to start the training")
|
|
print("Training phase: training the policy...")
|
|
task.train(signal_from_interface=False)
|
|
print("Training phase: policy trained!")
|
|
|
|
# plot what the DMP policy has learned by performing a rollout
|
|
policy.plot_rollout(nrows=3, ncols=3, suptitle='DMP position trajectories in joint space',
|
|
titles=['q' + str(i) for i in range(robot.num_actuated_joints)], show=True)
|
|
|
|
# test policy in simulation
|
|
print("Reproduction phase: test policy in simulation...")
|
|
# robot.set_home_joint_positions()
|
|
task.test(num_steps=rate*100, signal_from_interface=False)
|
|
print("Reproduction phase: Policy tested!")
|
|
|
|
# test policy on real robot
|
|
if use_real_robot:
|
|
input("Press Enter to move to the real robot experiment...")
|
|
sim.enable_middleware() # enable the real robot
|
|
ros.switch_mode(subscribe=False, publish=True, teleoperate=True)
|
|
print("Reproduction phase: test policy in reality...")
|
|
# robot.set_home_joint_positions()
|
|
task.test(num_steps=rate*100, signal_from_interface=False)
|
|
print("Reproduction phase: Policy tested!")
|