From 4c35bc6351986a1baa6c97ab7b8fcf99395a6a17 Mon Sep 17 00:00:00 2001
From: Younggyo Seo <seo0gyo@gmail.com>
Date: Tue, 2 Jul 2024 10:40:06 +0100
Subject: [PATCH] Support setting max_velocity and max_acceleration for arms
 (#237)

* support setting max_velocity and max_acceleration for arms

* fix a conflict in dataset_generator
---
 rlbench/dataset_generator.py |  4 ++++
 rlbench/environment.py       | 11 ++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/rlbench/dataset_generator.py b/rlbench/dataset_generator.py
index 4b06671d0..f05bbe2ae 100644
--- a/rlbench/dataset_generator.py
+++ b/rlbench/dataset_generator.py
@@ -192,6 +192,8 @@ def run(i, lock, task_index, variation_count, results, file_lock, tasks, args):
     rlbench_env = Environment(
         action_mode=MoveArmThenGripper(JointVelocity(), Discrete()),
         obs_config=obs_config,
+        arm_max_velocity=args.arm_max_velocity,
+        arm_max_acceleration=args.arm_max_acceleration,
         headless=True)
     rlbench_env.launch()
 
@@ -287,6 +289,8 @@ def parse_args():
     parser.add_argument('--processes', type=int, default=1, help='The number of parallel processes during collection.')
     parser.add_argument('--episodes_per_task', type=int, default=10, help='The number of episodes to collect per task.')
     parser.add_argument('--variations', type=int, default=-1, help='Number of variations to collect per task. -1 for all.')
+    parser.add_argument('--arm_max_velocity', type=float, default=1.0, help='Max arm velocity used for motion planning.')
+    parser.add_argument('--arm_max_acceleration', type=float, default=1.0, help='Max arm acceleration used for motion planning.')
     return parser.parse_args()
 
 
diff --git a/rlbench/environment.py b/rlbench/environment.py
index b1fb2f07f..4e9cb1a23 100644
--- a/rlbench/environment.py
+++ b/rlbench/environment.py
@@ -1,4 +1,5 @@
 import importlib
+from functools import partial
 from os.path import exists, dirname, abspath, join
 from typing import Type, List
 
@@ -38,7 +39,9 @@ def __init__(self,
                  visual_randomization_config: VisualRandomizationConfig = None,
                  dynamics_randomization_config: DynamicsRandomizationConfig = None,
                  attach_grasped_objects: bool = True,
-                 shaped_rewards: bool = False
+                 shaped_rewards: bool = False,
+                 arm_max_velocity: float = 1.0,
+                 arm_max_acceleration: float = 4.0,
                  ):
 
         self._dataset_root = dataset_root
@@ -54,6 +57,8 @@ def __init__(self,
         self._dynamics_randomization_config = dynamics_randomization_config
         self._attach_grasped_objects = attach_grasped_objects
         self._shaped_rewards = shaped_rewards
+        self._arm_max_velocity = arm_max_velocity
+        self._arm_max_acceleration = arm_max_acceleration
 
         if robot_setup not in SUPPORTED_ROBOTS.keys():
             raise ValueError('robot_configuration must be one of %s' %
@@ -97,6 +102,10 @@ def launch(self):
 
         arm_class, gripper_class, _ = SUPPORTED_ROBOTS[
             self._robot_setup]
+        arm_class = partial(
+            arm_class,
+            max_velocity=self._arm_max_velocity,
+            max_acceleration=self._arm_max_acceleration)
 
         # We assume the panda is already loaded in the scene.
         if self._robot_setup != 'panda':