From 1375cedcb5226e907463193f257259d8e99282be Mon Sep 17 00:00:00 2001
From: Michel Breyer <michel.p.breyer@gmail.com>
Date: Fri, 9 Jul 2021 15:53:34 +0200
Subject: [PATCH] Seed the rng of the simulation

---
 active_grasp/simulation.py | 23 ++++++++++++-----------
 config/active_grasp.yaml   |  4 ++++
 launch/active_grasp.launch |  2 +-
 scripts/bt_sim_node.py     | 21 +++++++++------------
 4 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/active_grasp/simulation.py b/active_grasp/simulation.py
index 9c9d85b..d1c7d72 100644
--- a/active_grasp/simulation.py
+++ b/active_grasp/simulation.py
@@ -4,25 +4,26 @@ import rospkg
 
 from active_grasp.utils import *
 from robot_utils.bullet import *
-from robot_utils.controllers import CartesianPoseController
 from robot_utils.spatial import Rotation, Transform
 
 
 class Simulation(BtSim):
-    def __init__(self, gui=True):
+    def __init__(self, gui, rng):
         super().__init__(gui=gui, sleep=False)
-        # p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
-        p.resetDebugVisualizerCamera(1.4, 50, -35, [0.0, 0.0, 0.6])
-
+        self.rng = rng
         self.object_uids = []
 
+        self.configure_visualizer()
         self.find_object_urdfs()
         self.load_table()
         self.load_robot()
         self.load_controller()
-
         self.reset()
 
+    def configure_visualizer(self):
+        # p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
+        p.resetDebugVisualizerCamera(1.4, 50, -35, [0.0, 0.0, 0.6])
+
     def find_object_urdfs(self):
         rospack = rospkg.RosPack()
         root = Path(rospack.get_path("vgn")) / "assets/urdfs/packed/test"
@@ -54,8 +55,8 @@ class Simulation(BtSim):
 
     def set_initial_arm_configuration(self):
         q = self.arm.configurations["ready"]
-        q[0] = np.deg2rad(np.random.uniform(-10, 10))
-        q[5] = np.deg2rad(np.random.uniform(90, 105))
+        q[0] = np.deg2rad(self.rng.uniform(-10, 10))
+        q[5] = np.deg2rad(self.rng.uniform(90, 105))
         for i, q_i in enumerate(q):
             p.resetJointState(self.arm.uid, i, q_i, 0)
         p.resetJointState(self.arm.uid, 9, 0.04, 0)
@@ -79,7 +80,7 @@ class Simulation(BtSim):
     def load_random_object_arrangement(self, attempts=10):
         origin = np.r_[self.origin[:2], 0.625]
         scale = 0.8
-        urdfs = np.random.choice(self.urdfs, 4)
+        urdfs = self.rng.choice(self.urdfs, 4)
         for urdf in urdfs:
             # Load the object
             uid = self.load_object(urdf, Rotation.identity(), [0.0, 0.0, 0.0], scale)
@@ -88,8 +89,8 @@ class Simulation(BtSim):
             state_id = p.saveState()
             for _ in range(attempts):
                 # Try to place the object without collision
-                ori = Rotation.from_rotvec([0.0, 0.0, np.random.uniform(0, 2 * np.pi)])
-                offset = np.r_[np.random.uniform(0.2, 0.8, 2) * self.length, z_offset]
+                ori = Rotation.from_rotvec([0.0, 0.0, self.rng.uniform(0, 2 * np.pi)])
+                offset = np.r_[self.rng.uniform(0.2, 0.8, 2) * self.length, z_offset]
                 p.resetBasePositionAndOrientation(uid, origin + offset, ori.as_quat())
                 self.step()
                 if not p.getContactPoints(uid):
diff --git a/config/active_grasp.yaml b/config/active_grasp.yaml
index 6ff8ce8..1bd7e0f 100644
--- a/config/active_grasp.yaml
+++ b/config/active_grasp.yaml
@@ -1,3 +1,7 @@
+bt_sim:
+  gui: True
+  seed: 12
+
 active_grasp:
   frame_id: task
   length: 0.3
diff --git a/launch/active_grasp.launch b/launch/active_grasp.launch
index 6d055b9..9783ad0 100644
--- a/launch/active_grasp.launch
+++ b/launch/active_grasp.launch
@@ -5,7 +5,7 @@
   <rosparam command="load" file="$(find active_grasp)/config/active_grasp.yaml" subst_value="true" />
 
   <param name="robot_description" command="$(find xacro)/xacro $(find active_grasp)/assets/urdfs/panda_arm_hand.urdf.xacro" />
-  <node pkg="active_grasp" type="bt_sim_node.py" name="simulation" args="--gui" output="screen" />
+  <node pkg="active_grasp" type="bt_sim_node.py" name="bt_sim" output="screen" />
   <node pkg="robot_state_publisher" type="robot_state_publisher" name="robot_state_publisher" />
 
   <node pkg="rviz" type="rviz" name="rviz" args="-d $(find active_grasp)/launch/active_grasp.rviz" if="$(arg launch_rviz)" />
diff --git a/scripts/bt_sim_node.py b/scripts/bt_sim_node.py
index fc46646..e132c5d 100755
--- a/scripts/bt_sim_node.py
+++ b/scripts/bt_sim_node.py
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 import actionlib
-import argparse
 import cv_bridge
 import franka_gripper.msg
 from geometry_msgs.msg import PoseStamped
@@ -17,8 +16,10 @@ from robot_utils.ros.conversions import *
 
 
 class BtSimNode:
-    def __init__(self, gui):
-        self.sim = Simulation(gui=gui)
+    def __init__(self):
+        gui = rospy.get_param("~gui", True)
+        rng = self.get_rng()
+        self.sim = Simulation(gui=gui, rng=rng)
         self.robot_state_interface = RobotStateInterface(self.sim.arm, self.sim.gripper)
         self.arm_interface = ArmInterface(self.sim.arm, self.sim.controller)
         self.gripper_interface = GripperInterface(self.sim.gripper)
@@ -29,6 +30,10 @@ class BtSimNode:
         self.advertise_services()
         self.broadcast_transforms()
 
+    def get_rng(self):
+        seed = rospy.get_param("~seed", None)
+        return np.random.default_rng(seed) if seed else np.random
+
     def advertise_services(self):
         rospy.Service("reset", Reset, self.reset)
 
@@ -166,17 +171,9 @@ class CameraInterface:
         self.depth_pub.publish(depth_msg)
 
 
-def create_parser():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--gui", action="store_true")
-    return parser
-
-
 def main():
     rospy.init_node("bt_sim")
-    parser = create_parser()
-    args, _ = parser.parse_known_args()
-    server = BtSimNode(args.gui)
+    server = BtSimNode()
     server.run()