Gather demo collection at once (#72)

scripts/demo/open_manipulator/open_manipulator_demo_collector.py

@@ -8,7 +8,7 @@ import json
 import random
 import threading
 import time
-from collections import OrderedDict
+from collections import defaultdict
 from math import pi, pow
 import numpy as np
 import rospy
@@ -59,7 +59,7 @@ class DemoCollector(object):
         self.T_goal = np.array(self.robot.forward(self.q))
         self.T_cur = np.array(self.robot.forward(self.q))
 
-        self.num_tar_demo = 10
+        self.num_tar_demo = 3
         self.num_cur_demo = 0
 
         self.control_start_time = (
@@ -116,10 +116,11 @@ class DemoCollector(object):
         3) If target is set, move to target.
         4) If robot pose is not initial pose, move to initial pose.
         5) If target is set and robot pose is not initial pose, take ros sleep.
-        6) If number of target demo is bigger than number of current demo, finish demo 
+        6) If number of target demo is bigger than number of current demo, finish demo
         collection.
         """
         self.r = rospy.Rate(100)
+        self.start_log()
         while not self.is_finished:
             if self.is_joint_states_cb is True:
                 if self.init is False:
@@ -129,8 +130,7 @@ class DemoCollector(object):
                         rospy.get_rostime().secs + rospy.get_rostime().nsecs * 10 ** -9
                     )
                     self.init = True
-                if self.is_set_new_target is False and self.is_init_pos:
-                    self.start_log()  # run 2
+                if self.is_set_new_target is False and self.is_init_pos:  # run 2
                     self.set_target()
                     self.T_init = np.array(self.robot.forward(self.q))
                     self.control_start_time = (
@@ -147,6 +147,9 @@ class DemoCollector(object):
             if self.num_cur_demo > self.num_tar_demo:
                 print("Demo Collection Finished!")  # run 6
                 self.is_finished = True
+        with open("../DemoCollection.json", "w") as f:
+            json.dump(self.data, f)
+        print("Demo file saved successfully")
         quit()
 
     def joint_states_cb(self, joint_states):
@@ -161,7 +164,8 @@ class DemoCollector(object):
 
     def start_log(self):
         """ Start logging in .txt format."""
-        self.f = open("../DemoEpisode" + str(self.num_cur_demo) + ".txt", "w")
+        self.data = defaultdict(lambda: defaultdict(list))
+#        self.f = open("../DemoEpisode" + str(self.num_cur_demo) + ".txt", "w")
 
     def set_target(self):
         """ Randomly set target within joint limit and workspace limit."""
@@ -258,7 +262,6 @@ class DemoCollector(object):
         self.j4_pos_command_pub.publish(self.q_desired[3])
         self.joint_pos_command_to_dxl_pub.publish(data=self.q_desired)
 
-        data = OrderedDict()
         obs = np.concatenate(
             (
                 self._gripper_pos,
@@ -268,10 +271,11 @@ class DemoCollector(object):
                 self.effort,
             )
         )
-        data["observation"] = obs.tolist()
-        data["desired q"] = self.q_desired.tolist()
-        data["target"] = self.T_target[0:3, 3].tolist()
-        json.dump(data, self.f, ensure_ascii=False)
+
+        # append data
+        self.data[self.num_cur_demo]["observation"].append(obs.tolist())
+        self.data[self.num_cur_demo]["desired q"].append(self.q_desired.tolist())
+        self.data[self.num_cur_demo]["target"].append(self.T_target[0:3, 3].tolist())
 
         if np.mean(np.abs(self.T_target[0:3, 3] - self.T_cur[0:3, 3])) < 0.001:
             self.is_set_new_target = False
@@ -280,7 +284,6 @@ class DemoCollector(object):
             self.control_start_time = (
                 rospy.get_rostime().secs + rospy.get_rostime().nsecs * 10 ** -9
             )
-            self.f.close()
             print("Target arrived!")
 
     def move_to_init(self):