Use id maps in behavior nodes to reduce searching. (#517)

ihmcrobotics · Dec 4, 2024 · 5051931 · 5051931
1 parent 00b5ccf
commit 5051931
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 32 deletions.
diff --git a/...-high-level-behaviors/src/libgdx/java/us/ihmc/rdx/ui/behavior/sequence/RDXActionNode.java b/...-high-level-behaviors/src/libgdx/java/us/ihmc/rdx/ui/behavior/sequence/RDXActionNode.java
@@ -109,8 +109,7 @@ else if (definition.getExecuteAfterBeginning().getValue())
          }
          else
          {
-            int executeAfterActionIndex = state.calculateExecuteAfterActionIndex(actionSequence.getActionChildren());
-            ActionNodeState<?> executeAfterAction = actionSequence.getActionChildren().get(executeAfterActionIndex);
+            ActionNodeState<?> executeAfterAction = state.findExecuteAfterAction();
             selectedText = executeAfterAction.getDefinition().getName();
          }
 

diff --git a/...level-behaviors/src/libgdx/java/us/ihmc/rdx/ui/behavior/tree/RDXBehaviorTreeRootNode.java b/...level-behaviors/src/libgdx/java/us/ihmc/rdx/ui/behavior/tree/RDXBehaviorTreeRootNode.java
@@ -1,5 +1,6 @@
 package us.ihmc.rdx.ui.behavior.tree;
 
+import gnu.trove.map.hash.TLongObjectHashMap;
 import imgui.ImGui;
 import us.ihmc.behaviors.behaviorTree.BehaviorTreeRootNodeDefinition;
 import us.ihmc.behaviors.behaviorTree.BehaviorTreeRootNodeState;
@@ -20,6 +21,7 @@ public class RDXBehaviorTreeRootNode extends RDXBehaviorTreeNode<BehaviorTreeRoo
    private final BehaviorTreeRootNodeState state;
    private final ImBooleanWrapper automaticExecutionCheckbox;
    private final ImBooleanWrapper concurrencyEnabledCheckbox;
+   private final TLongObjectHashMap<RDXBehaviorTreeNode<?, ?>> idToNodeMap = new TLongObjectHashMap<>();
    private final List<RDXActionNode<?, ?>> actionChildren = new ArrayList<>();
    private final List<RDXActionNode<?, ?>> nextForExecutionActions = new ArrayList<>();
    private final List<RDXActionNode<?, ?>> currentlyExecutingActions = new ArrayList<>();
@@ -44,6 +46,7 @@ public void update()
    {
       super.update();
 
+      idToNodeMap.clear();
       actionChildren.clear();
       nextForExecutionActions.clear();
       currentlyExecutingActions.clear();
@@ -57,6 +60,8 @@ public void update()
 
    public void updateActionSubtree(RDXBehaviorTreeNode<?, ?> node)
    {
+      idToNodeMap.put(node.getState().getID(), node);
+
       for (RDXBehaviorTreeNode<?, ?> child : node.getChildren())
       {
          if (child instanceof RDXActionNode<?, ?> actionNode)
@@ -72,10 +77,8 @@ public void updateActionSubtree(RDXBehaviorTreeNode<?, ?> node)
                currentlyExecutingActions.add(actionNode);
             }
          }
-         else
-         {
-            updateActionSubtree(child);
-         }
+
+         updateActionSubtree(child);
       }
    }
 
@@ -178,4 +181,9 @@ public void renderNodeSettingsWidgets()
 
       super.renderNodeSettingsWidgets();
    }
+
+   public TLongObjectHashMap<RDXBehaviorTreeNode<?, ?>> getIDToNodeMap()
+   {
+      return idToNodeMap;
+   }
 }
diff --git a/...-behaviors/src/main/java/us/ihmc/behaviors/behaviorTree/BehaviorTreeRootNodeExecutor.java b/...-behaviors/src/main/java/us/ihmc/behaviors/behaviorTree/BehaviorTreeRootNodeExecutor.java
@@ -1,5 +1,6 @@
 package us.ihmc.behaviors.behaviorTree;
 
+import gnu.trove.map.hash.TLongObjectHashMap;
 import us.ihmc.behaviors.sequence.ActionNodeExecutor;
 import us.ihmc.communication.crdt.CRDTInfo;
 import us.ihmc.tools.io.WorkspaceResourceDirectory;
@@ -11,6 +12,7 @@ public class BehaviorTreeRootNodeExecutor extends BehaviorTreeNodeExecutor<Behav
 {
    private final BehaviorTreeRootNodeState state;
    private final BehaviorTreeRootNodeDefinition definition;
+   private final TLongObjectHashMap<BehaviorTreeNodeExecutor<?, ?>> idToNodeMap = new TLongObjectHashMap<>();
    private final List<ActionNodeExecutor<?, ?>> actionChildren = new ArrayList<>();
    private final List<ActionNodeExecutor<?, ?>> currentlyExecutingActions = new ArrayList<>();
 
@@ -35,6 +37,7 @@ public void update()
    {
       super.update();
 
+      idToNodeMap.clear();
       actionChildren.clear();
       currentlyExecutingActions.clear();
       updateActionSubtree(this);
@@ -56,8 +59,8 @@ public void update()
          int j = i - 1;
          for (; j >= 0; j--)
          {
-            int thisExecuteAfterActionIndex = state.getActionChildren().get(i).calculateExecuteAfterActionIndex(getState().getActionChildren());
-            int executeAfterActionIndexToCompare = state.getActionChildren().get(j).calculateExecuteAfterActionIndex(getState().getActionChildren());
+            int thisExecuteAfterActionIndex = state.getActionChildren().get(i).calculateExecuteAfterActionIndex();
+            int executeAfterActionIndexToCompare = state.getActionChildren().get(j).calculateExecuteAfterActionIndex();
             if (thisExecuteAfterActionIndex == executeAfterActionIndexToCompare)
             {
                state.getActionChildren().get(i).setConcurrencyRank(2);
@@ -77,7 +80,7 @@ else if (i == executionNextIndex)
          {
             state.getActionChildren().get(i).setIsNextForExecution(true);
          }
-         else if (state.getActionChildren().get(i).calculateExecuteAfterActionIndex(state.getActionChildren()) < executionNextIndex)
+         else if (state.getActionChildren().get(i).calculateExecuteAfterActionIndex() < executionNextIndex)
          {
             state.getActionChildren().get(i).setIsNextForExecution(true);
          }
@@ -127,6 +130,8 @@ else if (state.pollManualExecutionRequested())
 
    public void updateActionSubtree(BehaviorTreeNodeExecutor<?, ?> node)
    {
+      idToNodeMap.put(node.getState().getID(), node);
+
       for (BehaviorTreeNodeExecutor<?, ?> child : node.getChildren())
       {
          if (child instanceof ActionNodeExecutor<?, ?> actionNode)
@@ -137,10 +142,8 @@ public void updateActionSubtree(BehaviorTreeNodeExecutor<?, ?> node)
                currentlyExecutingActions.add(actionNode);
             }
          }
-         else
-         {
-            updateActionSubtree(child);
-         }
+
+         updateActionSubtree(child);
       }
    }
 
@@ -175,7 +178,7 @@ private boolean shouldExecuteNextAction()
 
       if (state.getConcurrencyEnabled())
       {
-         int executeAfterActionIndex = nextNodeToExecute.getState().calculateExecuteAfterActionIndex(getState().getActionChildren());
+         int executeAfterActionIndex = nextNodeToExecute.getState().calculateExecuteAfterActionIndex();
 
          if (executeAfterActionIndex < 0) // Execute after beginning
          {
@@ -201,6 +204,11 @@ private boolean isEndOfSequence()
    {
       return state.getExecutionNextIndex() >= actionChildren.size();
    }
+
+   public TLongObjectHashMap<BehaviorTreeNodeExecutor<?, ?>> getIDToNodeMap()
+   {
+      return idToNodeMap;
+   }
 
    public List<ActionNodeExecutor<?, ?>> getActionChildren()
    {

diff --git a/...vel-behaviors/src/main/java/us/ihmc/behaviors/behaviorTree/BehaviorTreeRootNodeState.java b/...vel-behaviors/src/main/java/us/ihmc/behaviors/behaviorTree/BehaviorTreeRootNodeState.java
@@ -1,6 +1,7 @@
 package us.ihmc.behaviors.behaviorTree;
 
 import behavior_msgs.msg.dds.BehaviorTreeRootNodeStateMessage;
+import gnu.trove.map.hash.TLongObjectHashMap;
 import org.apache.commons.lang3.mutable.MutableInt;
 import us.ihmc.behaviors.sequence.ActionNodeState;
 import us.ihmc.communication.crdt.CRDTBidirectionalBoolean;
@@ -23,7 +24,8 @@ public class BehaviorTreeRootNodeState extends BehaviorTreeNodeState<BehaviorTre
    private final CRDTBidirectionalNotification manualExecutionRequested;
    private final CRDTBidirectionalBoolean concurrencyEnabled;
 
-   private transient final MutableInt actionIndex = new MutableInt();
+   private final TLongObjectHashMap<BehaviorTreeNodeState<?>> idToNodeMap = new TLongObjectHashMap<>();
+   private transient final MutableInt actionIndexAssignment = new MutableInt();
    private final List<ActionNodeState<?>> actionChildren = new ArrayList<>();
 
    public BehaviorTreeRootNodeState(long id, CRDTInfo crdtInfo, WorkspaceResourceDirectory saveFileDirectory)
@@ -43,24 +45,25 @@ public void update()
    {
       super.update();
 
-      actionIndex.setValue(0);
+      idToNodeMap.clear();
+      actionIndexAssignment.setValue(0);
       actionChildren.clear();
-      updateActionSubtree(this, actionIndex);
+      updateActionSubtree(this, actionIndexAssignment);
    }
 
    public void updateActionSubtree(BehaviorTreeNodeState<?> node, MutableInt actionIndex)
    {
+      idToNodeMap.put(node.getState().getID(), node);
+
       for (BehaviorTreeNodeState<?> child : node.getChildren())
       {
          if (child instanceof ActionNodeState<?> actionNode)
          {
             actionNode.setActionIndex(actionIndex.getAndIncrement());
             actionChildren.add(actionNode);
          }
-         else
-         {
-            updateActionSubtree(child, actionIndex);
-         }
+
+         updateActionSubtree(child, actionIndex);
       }
    }
 
@@ -166,6 +169,11 @@ public void setConcurrencyEnabled(boolean concurrencyEnabled)
       this.concurrencyEnabled.setValue(concurrencyEnabled);
    }
 
+   public TLongObjectHashMap<BehaviorTreeNodeState<?>> getIDToNodeMap()
+   {
+      return idToNodeMap;
+   }
+
    public List<ActionNodeState<?>> getActionChildren()
    {
       return actionChildren;

diff --git a/ihmc-high-level-behaviors/src/main/java/us/ihmc/behaviors/sequence/ActionNodeState.java b/ihmc-high-level-behaviors/src/main/java/us/ihmc/behaviors/sequence/ActionNodeState.java
@@ -2,6 +2,8 @@
 
 import behavior_msgs.msg.dds.ActionNodeStateMessage;
 import us.ihmc.behaviors.behaviorTree.BehaviorTreeNodeState;
+import us.ihmc.behaviors.behaviorTree.BehaviorTreeRootNodeState;
+import us.ihmc.behaviors.behaviorTree.BehaviorTreeTools;
 import us.ihmc.communication.crdt.CRDTInfo;
 import us.ihmc.communication.crdt.CRDTStatusDoubleArray;
 import us.ihmc.communication.crdt.CRDTStatusOneDoFJointTrajectoryList;
@@ -292,27 +294,34 @@ public boolean getIsExecuting()
       return isExecuting.getValue();
    }
 
-   public int calculateExecuteAfterActionIndex(List<ActionNodeState<?>> actionStateChildren)
+   public int calculateExecuteAfterActionIndex()
    {
       if (definition.getExecuteAfterBeginning().getValue())
       {
          return -1;
       }
       else if (!definition.getExecuteAfterPrevious().getValue())
       {
-         long executeAfterID = definition.getExecuteAfterNodeID().getValue();
-         for (int j = actionIndex - 1; j >= 0; j--)
-         {
-            ActionNodeState<?> action = actionStateChildren.get(j);
-            if (action.getID() == executeAfterID)
-            {
-               return action.getActionIndex();
-            }
-         }
+         ActionNodeState<?> executeAfterAction = findExecuteAfterAction();
 
-         LogTools.error("Action ID not found: {}", executeAfterID);
+         if (executeAfterAction != null)
+            return executeAfterAction.getActionIndex();
       }
 
       return actionIndex - 1; // previous
    }
+
+   public ActionNodeState<?> findExecuteAfterAction()
+   {
+      long executeAfterID = definition.getExecuteAfterNodeID().getValue();
+
+      if (BehaviorTreeTools.findRootNode(this).getIDToNodeMap().get(executeAfterID) instanceof ActionNodeState<?> executeAfterAction)
+      {
+         return executeAfterAction;
+      }
+
+      LogTools.error("Action ID not found: {}", executeAfterID);
+
+      return null;
+   }
 }