[Paper Notes] DEX: Demonstration-Guided Reinforcement Learning with Efficient Exploration for Task Automation of Surgical Robot

Contents

Part One: Paper Notes
1.1 Objective
1.2 Method
Part Two: Algorithm Reproduction
2.1 Environment Setup
2.2 Data Collection
2.3 Training
Part Three: Understanding the Code
3.1 Basic Definitions
Part Four: Code Modifications
4.1 Additional Environment Setup
4.2 Testing

Part One: Paper Notes

Title: Demonstration-Guided Reinforcement Learning with Efficient Exploration for Task Automation of Surgical Robot Title: Demonstration-Guided Reinforcement Learning with Efficient Exploration for Task Automation of Surgical Robot Author Team: The Chinese University of Hong Kong (Yun-Hui Liu’s team) Journal/Conference: ICRA Year: 2023 Code: https://github.com/med-air/DEX

1.1 Objective

Although reinforcement learning-based methods offer a possible approach to surgical automation, they typically require large amounts of collected data for learning. This paper therefore aims to improve the efficiency of learning through exploration from demonstrations and make effective use of expert demonstration data.

Specifically, the current problems are as follows:

With reinforcement learning, learning through exploration without demonstration data requires collecting a large amount of data to solve a task;
Methods that use demonstration data—for example, giving demonstration data higher priority than the robot’s exploration data—remain inefficient. Methods that introduce additional reward functions are not only limited to specific environments but are also prone to local optima;
Methods that use an actor-critic framework regularize the actor loss to measure the behavioral difference between the robot and the expert. However, this approach is inefficient, especially early on when the gap between the robot and the demonstrations is large. It also does not regularize the critic, which can easily lead to overestimation.

The contributions of this paper are:

It proposes an actor-critic framework that reduces critic overestimation and encourages exploration through expert-like actions during reinforcement learning.
It uses nonparametric guidance propagation to enable exploration of unobserved states.
Experiments on the SurRoL surgical robot demonstrate excellent performance. Deployment on the dVRK also shows strong potential.

dVRK (da Vinci Research Kit, da Vinci surgical robot system)

1.2 Method

DEX (Demonstration-guided EXploration), or demonstration-guided exploration.

(0) Problem Definition

Surgical robot action learning is formulated as an off-policy learning problem in which an agent interacts with an environment constructed as a Markov decision process.

off-policy means that the agent does not use its current policy to select actions. Instead, it uses a different policy to generate behavioral data and learns an optimal method for making behavioral decisions from past experience.

At time $t$ , the robot executes an action based on the current state $s_t$ and deterministic policy $\pi$ . The environment rewards the agent with $r_t=r(s_t,a_t)$ , and the state then transitions to $s_{t+1}$ .

This process repeats, and each time the agent stores the experience $(s+t,a_t,r_t,s_{t+1})$ in the replay buffer $D_A$ .

A demonstration buffer $D_E$ is also created to store experience from the expert policy $\pi$ .

Method

As shown in the figure, the method consists of two parts:

An actor-critic-based policy learning module (bottom right), which uses demonstration data to guide exploration;
A nonparametric module based on nearest-neighbor matching and locally weighted regression (top left), which propagates demonstrations that differ substantially from the current state for use at the current state.

(1) Expert-Guided Actor-Critic Framework

Existing actor-critic methods learn an optimal policy by maximizing expected returns, but inaccurate Q-value estimates can hinder exploration. This paper uses the action gap between the agent policy and the expert policy to augment the environment reward.

\max_{\pi}\mathbb{E}_{\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}(r_{t}-\alpha d(a_{t},a_{t}^{e}))\right],a_{t}^{e}:=\pi^{e}(s_{t}),

Here, $\alpha$ is the exploration coefficient, and $d()$ is a distance metric that measures the similarity between the agent action and the expert action.

Based on this reward, the paper designs a regularized Q-function (critic) and minimizes the gap between the action value and the state value.

(2) Guidance Propagation with Limited Demonstrations

During the initial learning stage, the agent can easily explore regions not covered by the demonstrations, making it impossible to supervise the actor’s exploration.

A conventional solution is behavior cloning, but when states differ greatly, the policy’s actions can still differ substantially from the expert actions. The paper therefore uses a nonparametric regression model to propagate experience from limited demonstrations and provide more stable guidance.

First, a small batch of states and actions is sampled from the demonstration buffer. Then, given a current state, the method searches the small batch and uses k-nearest neighbors to find the closest states, after which a locally weighted regression method using an exponential sum function approximates the expert policy.

\hat{\pi}^e(s)=\frac{\sum_{i=1}^k\exp\left(-\|s-s^{(i)}\|_2\right)\cdot a^{(i)}}{\sum_{i=1}^k\exp\left(-\|s-s^{(i)}\|_2)\right)}.

Part Two: Algorithm Reproduction

2.1 Environment Setup

Clone the code

git clone --recursive https://github.com/med-air/DEX.git
cd DEX

Create a virtual environment

conda create -n dex python=3.8
conda activate dex

Install the dependencies

pip3 install -e SurRoL/	# install surrol environments
pip3 install -r requirements.txt
pip3 install -e .

# directory: anaconda3/envs/dex/lib/python3.8/site-packages/gym/envs/__init__.py
import surrol.gym

2.2 Data Collection

mkdir SurRoL/surrol/data/demo
python SurRoL/surrol/data/data_generation.py --env NeedlePick-v0

2.3 Training

python3 train.py task=NeedlePick-v0 agent=dex use_wb=True

The program provided by the authors also includes other reinforcement learning and imitation learning algorithms, such as:

DDPG: deep deterministic policy gradient reinforcement learning
DDPGBC: deep deterministic policy gradient reinforcement learning + behavior cloning
SAC: maximum-entropy model-free deep reinforcement learning
SQIL: imitation learning with regularized behavior cloning
COL: behavior cloning and reinforcement learning
AWAC: offline reinforcement learning
AMP: adversarial imitation learning

Part Three: Understanding the Code

3.1 Basic Definitions

(1) Robot State

The robot state is obtained with the following function.

# SurRoL.surrol.tasks.psm_env.PsmEnv._get_robot_state
def _get_robot_state(self, idx: int) -> np.ndarray:
	'''获取机器人的状态，返回机器人当前位姿、夹爪角度，两者拼接成一个数组（3位置+3欧拉角+1开合角度）'''
	# robot state: tip pose in the world coordinate
	psm = self.psm1 if idx == 0 else self.psm2
	pose_world = psm.pose_rcm2world(psm.get_current_position(), 'tuple')    # 机器人在世界坐标系下的位姿
	jaw_angle = psm.get_current_jaw_position()  # 夹爪角度
	return np.concatenate([
		np.array(pose_world[0]), np.array(p.getEulerFromQuaternion(pose_world[1])), np.array(jaw_angle).ravel()
	])  # 3 + 3 + 1 = 7

The robot state is therefore an array containing the robot’s current pose and gripper angle: three position values, three Euler angles, and one opening angle.

robot_state = [x, y, z, roll, pitch, yaw, gripper]

(2) Observation State

The observation state is obtained with the following method, which gathers the robot’s current position, the target object’s position, and their relative position:

# SurRoL.surrol.tasks.psm_env.PsmEnv._get_obs
def _get_obs(self) -> dict:
	'''获取当前环境状态信息（机器人当前位置、目标物体位置、机器人与目标物体的相对位置）'''
	# 获取机器人当前状态
	robot_state = self._get_robot_state(idx=0)
	# TODO: may need to modify
	
	# 检查环境中是否有物体，如果有则获取目标物体的位置、姿态、相对于机器人的位置
	if self.has_object:
		pos, _ = get_link_pose(self.obj_id, -1)     # 目标物体位置
		object_pos = np.array(pos)
		pos, orn = get_link_pose(self.obj_id, self.obj_link1)   # 获取目标物体的特定link的位置和方向
		waypoint_pos = np.array(pos)    # 路径点位置为目标物体link位置
		# rotations
		waypoint_rot = np.array(p.getEulerFromQuaternion(orn))  #路径点姿态为目标物体link姿态
		# relative position state
		object_rel_pos = object_pos - robot_state[0: 3] # 相对位置为目标位置与机器人末端位置之差
	else:
		# TODO: can have a same-length state representation
		object_pos = waypoint_pos = waypoint_rot = object_rel_pos = np.zeros(0)

	# 确定使用哪个位置作为目标位置
	if self.has_object:
		# object/waypoint position，使用物体位置object_pos，或物体link位置waypoint_pos
		achieved_goal = object_pos.copy() if not self._waypoint_goal else waypoint_pos.copy()
	else:
		# tip position，如果没有目标物体，则将机器人末端的位置作为目标位置
		achieved_goal = np.array(get_link_pose(self.psm1.body, self.psm1.TIP_LINK_INDEX)[0])

	observation = np.concatenate([
		robot_state, object_pos.ravel(), object_rel_pos.ravel(),
		waypoint_pos.ravel(), waypoint_rot.ravel()  # achieved_goal.copy(),
	])
	obs = {
		'observation': observation.copy(),
		'achieved_goal': achieved_goal.copy(),
		'desired_goal': self.goal.copy()
	}
	return obs

As shown, the observation state is a dictionary containing three keys: observation, achieved_goal, and desired_goal.

observation: consists of robot_state, object_pos.ravel(), object_rel_pos.ravel(), waypoint_pos.ravel(), and waypoint_rot.ravel().

obs	Key value (all three are one-dimensional arrays)	Meaning
observation	[robot_x, robot_y, robot_z, robot_roll, robot_pitch, robot_yaw, gripper,	Robot end-effector pose and gripper state
	obj_x, obj_y, obj_z,	Target object position
	obj_rel_x, obj_rel_y, obj_rel_z,	Position of the target object relative to the current robotic arm
	obj_link1_x, obj_link1_y, obj_link1_z	Position of the target object’s link1 used as the waypoint position
	obj_link1_roll, obj_link1_pitch, obj_link1_yaw]	Pose of the target object’s link1 used as the waypoint pose
achieved_goal	[x, y, z]	The waypoint is used as the actual position; if link1 is not set, the target object position is used; if there is no object, the robot end-effector position is used
desired_goal	[x, y, z]	Target position of the robot end effector

(3) Robot Action

Robot actions are executed with the following method, which applies changes in position and rotation and opens or closes the gripper:

# SurRoL.surrol.tasks.psm_env.PsmEnv._set_action
def _set_action(self, action: np.ndarray):
	"""
	delta_position (3), delta_theta (1) and open/close the gripper (1)
	in the world frame
	执行动作的过程（位置的变化，旋转变化，夹爪的开合）
	"""
	assert len(action) == self.ACTION_SIZE, "The action should have the save dim with the ACTION_SIZE" # ACTION_SIZE = 5
	# time0 = time.time()
	action = action.copy()  # ensure that we don't change the action outside of this scope
	action[:3] *= 0.01 * self.SCALING  # position, limit maximum change in position
	pose_world = self.psm1.pose_rcm2world(self.psm1.get_current_position())
	workspace_limits = self.workspace_limits1
	pose_world[:3, 3] = np.clip(pose_world[:3, 3] + action[:3],
								workspace_limits[:, 0] - [0.02, 0.02, 0.],
								workspace_limits[:, 1] + [0.02, 0.02, 0.08])  # clip to ensure convergence
	rot = get_euler_from_matrix(pose_world[:3, :3])
	if self.ACTION_MODE == 'yaw':
		action[3] *= np.deg2rad(30)  # yaw, limit maximum change in rotation
		rot = (self.psm1_eul[0], self.psm1_eul[1], wrap_angle(rot[2] + action[3]))  # only change yaw
	elif self.ACTION_MODE == 'pitch':
		action[3] *= np.deg2rad(15)  # pitch, limit maximum change in rotation
		pitch = np.clip(wrap_angle(rot[1] + action[3]), np.deg2rad(-90), np.deg2rad(90))
		rot = (self.psm1_eul[0], pitch, self.psm1_eul[2])  # only change pitch
	else:
		raise NotImplementedError
	pose_world[:3, :3] = get_matrix_from_euler(rot)
	action_rcm = self.psm1.pose_world2rcm(pose_world)
	# time1 = time.time()
	self.psm1.move(action_rcm)
	# time2 = time.time()

	# jaw
	if self.block_gripper:
		action[4] = -1
	if action[4] < 0:
		self.psm1.close_jaw()
		self._activate(0)
	else:
		self.psm1.move_jaw(np.deg2rad(40))  # open jaw angle; can tune
		self._release(0)

The robot action is an array of length 5:

action[5] = [delta_x, delta_y, delta_z, yaw/pitch, gripper]

delta_xyz represents the displacement of the robot end effector along the x, y, and z axes;
yaw/pitch determines which rotation to use according to the class variable ACTION_MODE;
A negative gripper value closes the gripper; a nonnegative value determines how far it opens.

Part Four: Code Modifications

4.1 Additional Environment Setup

pip install hydra-core colorlog termcolor opencv-python perlin_noise
pip install "cython<3"

4.2 Testing

AWAC: offline reinforcement learning

(1) DDPGBC: Deep Deterministic Policy Gradient Reinforcement Learning + Behavior Cloning

DDPGBC	Demo 1000	Demo 10000
Success Reward	Failed	Failed
Trajectory Reward	Failed

(2) SAC: Maximum-Entropy Model-Free Deep Reinforcement Learning

SAC	Demo 1000	Demo 10000
Success Reward	Moved the object to the target position, with a high success rate	Moved the object to the target position, with a high success rate
Trajectory Reward	Moved the object to the target position, with a high success rate

(3) SQIL: Imitation Learning with Regularized Behavior Cloning

SQIL	Demo 1000	Demo 10000
Success Reward	Failed	Failed
Trajectory Reward	Failed

(4) COL: Behavior Cloning and Reinforcement Learning

CoL	Demo 1000	Demo 10000
Success Reward	Failed	Failed
Trajectory Reward	Failed

(5) AWAC: Offline Reinforcement Learning

AWAC	Demo 1000	Demo 10000
Success Reward	Moved the object to the target position, with a low success rate	Moved the object to the target position, with a high success rate
Trajectory Reward	Failed

(6) AMP: Adversarial Imitation Learning

AMP	Demo 1000	Demo 10000
Success Reward	Moved the object to the target position, with a high success rate	Moved the object to the target position, with a high success rate
Trajectory Reward	Moved the object to the target position, with a high success rate

(7) DEX: Demonstration-Guided Reinforcement Learning

DEX	Demo 1000	Demo 10000
Success Reward	Failed	Failed
Trajectory Reward	Failed