Separate timings for HER and rb [no ci]

Author: Forest Agostinelli

CHANGES.md

@@ -27,6 +27,7 @@
 * Add layer norm to resnet_fc
 * Vectorize expand
 * Make policy at update_num=0 sampled from Domain's sample_state_action
+* Separate timings for HER and rb
 
 ## 0.2.1
 * Consolidate search: Beam search -> special cases: greedy_policy, graph search -> special cases: batch weighted A* search, batch weighted Q* search 

deepxube/_solve.py

@@ -60,7 +60,7 @@ def parse_solve(parser: ArgumentParser) -> None:
     parser.add_argument('--pathfind', type=str, required=True, help="Pathfinding algorithm and arguments.")
     parser.add_argument('--file', type=str, required=True, help="File containing problem instances to solve")
 
-    parser.add_argument('--time_limit', type=float, default=-1.0, help="A time limit for search. Default is -1, which means infinite.")
+    parser.add_argument('--time_limit', type=float, default=-1.0, help="A time limit (in seconds) for search. Default is -1, which means infinite.")
     parser.add_argument('--max_itrs', type=int, default=None, help="Maximum number of search iterations. None for infinite.")
 
     parser.add_argument('--results', type=str, required=True, help="Directory to save results. Saves results after every instance.")
@@ -124,10 +124,7 @@ def solve_cli(args: argparse.Namespace) -> None:
     # heur and policy fn
     heur_fn: Optional[HeurFn] = get_heur_fn(domain, domain_name, args.heur, args.heur_file, args.heur_type, args.nnet_batch_size)
     policy_fn: Optional[PolicyFn] = get_policy_fn(domain, domain_name, args.policy, args.policy_file, args.policy_samp, args.nnet_batch_size)
-    print(domain)
     pathfind_functions: Any = get_pathfind_functions(get_pathfind_name_kwargs(args.pathfind)[0], heur_fn, policy_fn)
-    pathfind: PathFind = get_pathfind_from_arg(domain, pathfind_functions, args.pathfind)[0]
-    print(pathfind)
 
     # get data
     data: Dict = pickle.load(open(args.file, "rb"))
@@ -152,6 +149,11 @@ def solve_cli(args: argparse.Namespace) -> None:
         if not args.debug:
             sys.stdout = data_utils.Logger(output_file, "w")
 
+    # print info
+    print(domain)
+    pathfind: PathFind = get_pathfind_from_arg(domain, pathfind_functions, args.pathfind)[0]
+    print(pathfind)
+
     start_idx: int
     if args.start_idx is not None:
         start_idx = args.start_idx

deepxube/base/updater.py

@@ -511,12 +511,12 @@ def _get_her_goals(self, instances: List[Inst], times: Times) -> Tuple[List[Inst
                         state_deepest = node.state
                 states_deepest.append(state_deepest)
 
-            times.record_time("her_node_deepest", time.time() - start_time)
+            times.record_time("node_deepest", time.time() - start_time, path=["HER"])
 
             # relabel
             start_time = time.time()
             goals_relabel = self.domain.sample_goal_from_state(states_start, states_deepest)
-            times.record_time("her_relabel", time.time() - start_time)
+            times.record_time("relabel", time.time() - start_time, path=["HER"])
 
         return instances_goalkeep + instances_relabel, goals_goalkeep + goals_relabel
 

deepxube/updaters/updater_policy_rl.py

@@ -70,13 +70,13 @@ def _init_replay_buffer(self, max_size: int) -> None:
     def _rb_add(self, states: List[State], goals: List[Goal], actions: List[Action], times: Times) -> None:
         start_time = time.time()
         self.rb.add(list(zip(states, goals, actions, strict=True)))
-        times.record_time("rb_add", time.time() - start_time)
+        times.record_time("add", time.time() - start_time, path=["replay"])
 
     def _sample_rb(self, num: int, times: Times) -> Tuple[List[State], List[Goal], List[Action]]:
         # sample from replay buffer
         start_time = time.time()
         states, goals, actions = self.rb.sample(num)
-        times.record_time("rb_samp", time.time() - start_time)
+        times.record_time("samp", time.time() - start_time, path=["replay"])
 
         return states, goals, actions
 
@@ -155,7 +155,7 @@ def _get_instance_data_rb(self, instances: List[Instance], times: Times) -> List
             goals_her.extend([goal_her] * len(states_inst))
             actions_her.extend([edge.action for edge in instance.get_edges_popped()])
 
-        times.record_time("data_her", time.time() - start_time)
+        times.record_time("data", time.time() - start_time, path=["HER"])
 
         # add to replay buffer
         self._rb_add(states_her, goals_her, actions_her, times)

deepxube/updaters/updater_q_rl.py

@@ -56,8 +56,7 @@ def __init__(self, domain: D, pathfind_arg: str, up_args: UpArgs):
     def _step(self, pathfind: PathFindSetHeurQ, times: Times) -> None:
         _pathfind_q_step(pathfind)
 
-    def _q_learning_target(self, goals: List[Goal], is_solved_l: List[bool], tcs: List[float], states_next: List[State], times: Times) -> List[float]:
-        start_time = time.time()
+    def _q_learning_target(self, goals: List[Goal], is_solved_l: List[bool], tcs: List[float], states_next: List[State]) -> List[float]:
         # min cost-to-go for next state
         actions_next: List[List[Action]] = self.get_pathfind().get_state_actions(states_next, goals)
         qvals_next_l: List[List[float]] = self._get_targ_heur_fn()(states_next, goals, actions_next)
@@ -67,8 +66,6 @@ def _q_learning_target(self, goals: List[Goal], is_solved_l: List[bool], tcs: Li
         ctg_backups: NDArray = np.array(tcs) + np.array(qvals_next_min)
         ctg_backups = ctg_backups * np.logical_not(np.array(is_solved_l))
 
-        times.record_time("qlearn_targ", time.time() - start_time)
-
         return cast(List[float], ctg_backups.tolist())
 
     def _inputs_ctgs_to_np(self, states: List[State], goals: List[Goal], actions: List[Action], ctgs_backup: List[float], times: Times) -> List[NDArray]:
@@ -85,16 +82,18 @@ def _rb_add(self, states: List[State], goals: List[Goal], is_solved_l: List[bool
                 times: Times) -> None:
         start_time = time.time()
         self.rb.add(list(zip(states, goals, is_solved_l, actions, tcs, states_next, strict=True)))
-        times.record_time("rb_add", time.time() - start_time)
+        times.record_time("add", time.time() - start_time, path=["replay"])
 
     def _sample_rb_qlearn_target(self, num: int, times: Times) -> Tuple[List[State], List[Goal], List[Action], List[float]]:
         # sample from replay buffer
         start_time = time.time()
         states, goals, is_solved_l, actions, tcs, states_next = self.rb.sample(num)
-        times.record_time("rb_samp", time.time() - start_time)
+        times.record_time("samp", time.time() - start_time, path=["replay"])
 
         # value iteration update
-        ctgs_backup: List[float] = self._q_learning_target(goals, is_solved_l, tcs, states_next, times)
+        start_time = time.time()
+        ctgs_backup: List[float] = self._q_learning_target(goals, is_solved_l, tcs, states_next)
+        times.record_time("qlearn_targ", time.time() - start_time, path=["replay"])
 
         return states, goals, actions, ctgs_backup
 
@@ -201,12 +200,12 @@ def _get_instance_data_rb(self, instances: List[InstanceEdge], times: Times) ->
                 tcs_her.append(tc)
                 states_next_her.append(node_next.state)
 
-        times.record_time("data_her", time.time() - start_time)
+        times.record_time("data", time.time() - start_time, path=["HER"])
 
         # is solved
         start_time = time.time()
         is_solved_l_her: List[bool] = self.domain.is_solved(states_her, goals_her)
-        times.record_time("is_solved_her", time.time() - start_time)
+        times.record_time("is_solved", time.time() - start_time, path=["HER"])
 
         # add to replay buffer
         self._rb_add(states_her, goals_her, is_solved_l_her, actions_her, tcs_her, states_next_her, times)

deepxube/updaters/updater_v_rl.py

@@ -50,9 +50,7 @@ def __init__(self, domain: D, pathfind_arg: str, up_args: UpArgs):
     def _step(self, pathfind: PathFindSetHeurV, times: Times) -> None:
         _pathfind_v_step(pathfind)
 
-    def _value_iteration_target(self, goals: List[Goal], is_solved_l: List[bool], tcs_l: List[List[float]], states_exp: List[List[State]],
-                                times: Times) -> List[float]:
-        start_time = time.time()
+    def _value_iteration_target(self, goals: List[Goal], is_solved_l: List[bool], tcs_l: List[List[float]], states_exp: List[List[State]]) -> List[float]:
         # get cost-to-go of expanded states
         states_exp_flat, split_idxs = misc_utils.flatten(states_exp)
         goals_flat: List[Goal] = []
@@ -67,8 +65,6 @@ def _value_iteration_target(self, goals: List[Goal], is_solved_l: List[bool], tc
         ctgs_backup = np.array([np.min(x) for x in ctg_next_p_tc_l]) * np.logical_not(is_solved_l)
         ctgs_backup_l: List[float] = cast(List[float], ctgs_backup.tolist())
 
-        times.record_time("vi_targ", time.time() - start_time)
-
         return ctgs_backup_l
 
     def _inputs_ctgs_to_np(self, states: List[State], goals: List[Goal], ctgs_backup: List[float], times: Times) -> List[NDArray]:
@@ -85,21 +81,23 @@ def _init_replay_buffer(self, max_size: int) -> None:
     def _rb_add(self, states: List[State], goals: List[Goal], is_solved_l: List[bool], times: Times) -> None:
         start_time = time.time()
         self.rb.add(list(zip(states, goals, is_solved_l, strict=True)))
-        times.record_time("rb_add", time.time() - start_time)
+        times.record_time("add", time.time() - start_time, path=["replay"])
 
     def _sample_rb_vi_target(self, num: int, times: Times) -> Tuple[List[State], List[Goal], List[float]]:
         # sample from replay buffer
         start_time = time.time()
         states, goals, is_solved_l = self.rb.sample(num)
-        times.record_time("rb_samp", time.time() - start_time)
+        times.record_time("samp", time.time() - start_time, path=["replay"])
 
         # expand states
         start_time = time.time()
         states_exp, _, tcs_l = self.get_pathfind().expand_states(states, goals)
-        times.record_time("vi_expand", time.time() - start_time)
+        times.record_time("vi_expand", time.time() - start_time, path=["replay"])
 
         # value iteration update
-        ctgs_backup: List[float] = self._value_iteration_target(goals, is_solved_l, tcs_l, states_exp, times)
+        start_time = time.time()
+        ctgs_backup: List[float] = self._value_iteration_target(goals, is_solved_l, tcs_l, states_exp)
+        times.record_time("vi_targ", time.time() - start_time, path=["replay"])
 
         return states, goals, ctgs_backup
 
@@ -191,12 +189,12 @@ def _get_instance_data_rb(self, instances: List[InstanceNode], times: Times) ->
             states_her.extend(states_inst)
             goals_her.extend([goal_her] * len(states_inst))
 
-        times.record_time("data_her", time.time() - start_time)
+        times.record_time("data", time.time() - start_time, path=["HER"])
 
         # is solved
         start_time = time.time()
         is_solved_l_her: List[bool] = self.domain.is_solved(states_her, goals_her)
-        times.record_time("is_solved_her", time.time() - start_time)
+        times.record_time("is_solved", time.time() - start_time, path=["HER"])
 
         # add to replay buffer
         self._rb_add(states_her, goals_her, is_solved_l_her, times)