diff --git a/rdagent/log/ui/utils.py b/rdagent/log/ui/utils.py
index ade9833c..bdd5b053 100644
--- a/rdagent/log/ui/utils.py
+++ b/rdagent/log/ui/utils.py
@@ -612,16 +612,18 @@ def trace_figure(trace: Trace):
         """
         if hasattr(trace, "idx2loop_id"):
             # FIXME: only keep me after it is stable. Just for compatibility.
-            return f"L{trace.idx2loop_id[idx]}"
+            return f"L{trace.idx2loop_id[idx]} ({idx})"
         return f"L{idx}"
 
     # Add nodes and edges
     edges = []
+    parents_record = {}
     for i, parents in enumerate(trace.dag_parent):
         for parent in parents:
             edges.append((get_display_name(parent), get_display_name(i)))
         if len(parents) == 0:
             G.add_node(get_display_name(i))
+        parents_record[get_display_name(i)] = [get_display_name(parent) for parent in parents]
     G.add_edges_from(edges)
 
     # Check if G is a path (a single line)
@@ -658,27 +660,21 @@ def trace_figure(trace: Trace):
         pos = {}
 
         def parent_avg_pos(node):
-            id = int(node[1:])
-            parents = trace.dag_parent[id]
-
-            if not parents:
-                return 0
-
-            parent_nodes = [f"L{p}" for p in parents]
+            parent_nodes = parents_record.get(node, [])
             parent_xs = [pos[p][0] for p in parent_nodes if p in pos]
             return sum(parent_xs) / len(parent_xs) if parent_xs else 0
 
         for lvl in sorted(layer_nodes):
             nodes = layer_nodes[lvl]
             # For root nodes, sort directly by index
-            if lvl == 0:
-                sorted_nodes = sorted(nodes, key=lambda n: int(n[1:]))
+            if lvl == min(layer_nodes):
+                sorted_nodes = sorted(nodes, key=lambda n: int(n[1:].split(" ")[0]))
             else:
                 # Sort by average parent x, so children are below their parents
                 sorted_nodes = sorted(nodes, key=parent_avg_pos)
             y = -lvl  # y decreases as level increases (children below parents)
             for i, node in enumerate(sorted_nodes):
-                if lvl == 0:
+                if lvl == min(layer_nodes):
                     x = i
                 else:
                     # Place child directly below average parent x, offset if multiple at same y