From 2e05c7fe48a73df8b36abbfc717d5333b1aaf130 Mon Sep 17 00:00:00 2001
From: Whi Kwon <30768596+whikwon@users.noreply.github.com>
Date: Fri, 5 Jul 2019 18:40:47 +0900
Subject: [PATCH] Gather demo collection at once (#72)

---
 .../open_manipulator_demo_collector.py        | 27 ++++++++++---------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/scripts/demo/open_manipulator/open_manipulator_demo_collector.py b/scripts/demo/open_manipulator/open_manipulator_demo_collector.py
index 02637ac..1ef1937 100755
--- a/scripts/demo/open_manipulator/open_manipulator_demo_collector.py
+++ b/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):