Using Graph Structures to Represent Node State in Deep Reinforcement Learning (RL)-Based Decision Tree Construction
In one set of embodiments, a deep reinforcement learning (RL) system can train an agent to construct an efficient decision tree for classifying network packets according to a rule set, where the training includes: identifying, by an environment of the deep RL system, a leaf node in a decision tree; computing, by the environment, a graph structure representing a state of the leaf node, the graph structure including information regarding how one or more rules in the rule set that are contained in the leaf node are distributed in a hypercube of the leaf node; communicating, by the environment, the graph structure to the agent; providing, by the agent, the graph structure as input to a graph neural network; and generating, by the graph neural network based on the graph structure, an action to be taken on the leaf node for extending the decision tree.
1 . A method comprising:
training, by a deep reinforcement learning (RL) system, an agent to construct an efficient decision tree for classifying network packets according to a rule set comprising a plurality of rules, the training including:
identifying, by an environment of the deep RL system, a leaf node in a decision tree;
computing, by the environment, a graph structure representing a state of the leaf node, wherein the graph structure includes information regarding how one or more rules in the rule set that are contained in the leaf node are distributed in a hypercube of the leaf node, wherein the hypercube is a convex shape in a multi-dimensional space having dimensions corresponding to rule fields in the rule set, and wherein the hypercube's boundaries are defined by split conditions used to reach the leaf node in the decision tree;
communicating, by the environment, the graph structure to the agent;
providing, by the agent, the graph structure as input to a graph neural network, the graph neural network being designed to accept variable-sized graph structures as input; and
generating, by the graph neural network based on the graph structure, an action to be taken on the leaf node for extending the decision tree.
2 . The method of claim 1 wherein the information regarding how the one or more rules contained in the leaf node are distributed in the hypercube of the leaf node comprises, for each of the one or more rules, information regarding how a hypercube of said each rule is placed in the hypercube of the leaf node, the hypercube of said each rule being a convex shape in the multi-dimensional space having boundaries defined by matching patterns included in said each rule.
3 . The method of claim 1 wherein providing, by the agent, the graph structure as input to a graph neural network comprises:
transforming, by the agent, the graph structure using a graph transformation function; and
providing, by the agent, the transformed graph structure as input to the graph neural network.
4 . The method of claim 1 wherein the action splits the leaf node into a plurality of child nodes in accordance with one or more split conditions defined along one of the rule fields.
5 . The method of claim 1 wherein the graph structure is a bipartite graph comprising first and second sides, wherein the first side includes a plurality of graph nodes corresponding the plurality of rules, wherein the second side includes a grid of points in the hypercube of the leaf node, and wherein the bipartite graph includes an edge between a graph node in the plurality of graph nodes and a point in the grid of points if the point lies within a hypercube of the rule corresponding to the graph node.
6 . The method of claim 1 wherein the graph structure is a bipartite graph comprising first and second sides, wherein the first side includes a plurality of graph nodes corresponding the plurality of rules, wherein the second side includes a plurality of range trees corresponding to the rule fields of the rule set, and wherein the bipartite graph includes an edge between a graph node in the plurality of graph nodes and a node in a range tree corresponding to a rule field R if a matching pattern specified for rule field R in the rule corresponding to the graph node falls within a value range associated with the node.
7 . The method of claim 1 wherein the graph structure comprises a plurality of range trees corresponding to the rule fields of the rule set, and wherein at least one node of at least one range tree includes a number of rules in the rule set that match a value range associated with the node.
8 . A non-transitory computer readable storage medium having stored thereon program code executable by a computer system implementing a deep reinforcement learning (RL) system, the program code causing the computer system to execute a method comprising:
training an agent to construct an efficient decision tree for classifying network packets according to a rule set comprising a plurality of rules, the training including:
identifying, by an environment of the deep RL system, a leaf node in a decision tree;
computing, by the environment, a graph structure representing a state of the leaf node, wherein the graph structure includes information regarding how one or more rules in the rule set that are contained in the leaf node are distributed in a hypercube of the leaf node, wherein the hypercube is a convex shape in a multi-dimensional space having dimensions corresponding to rule fields in the rule set, and wherein the hypercube's boundaries are defined by split conditions used to reach the leaf node in the decision tree;
communicating, by the environment, the graph structure to the agent;
providing, by the agent, the graph structure as input to a graph neural network, the graph neural network being designed to accept variable-sized graph structures as input; and
generating, by the graph neural network based on the graph structure, an action to be taken on the leaf node for extending the decision tree.
9 . The non-transitory computer readable storage medium of claim 8 wherein the information regarding how the one or more rules contained in the leaf node are distributed in the hypercube of the leaf node comprises, for each of the one or more rules, information regarding how a hypercube of said each rule is placed in the hypercube of the leaf node, the hypercube of said each rule being a convex shape in the multi-dimensional space having boundaries defined by matching patterns included in said each rule.
10 . The non-transitory computer readable storage medium of claim 8 wherein providing, by the agent, the graph structure as input to a graph neural network comprises:
transforming, by the agent, the graph structure using a graph transformation function; and
providing, by the agent, the transformed graph structure as input to the graph neural network.
11 . The non-transitory computer readable storage medium of claim 8 wherein the action splits the leaf node into a plurality of child nodes in accordance with one or more split conditions defined along one of the rule fields.
12 . The non-transitory computer readable storage medium of claim 8 wherein the graph structure is a bipartite graph comprising first and second sides, wherein the first side includes a plurality of graph nodes corresponding the plurality of rules, wherein the second side includes a grid of points in the hypercube of the leaf node, and wherein the bipartite graph includes an edge between a graph node in the plurality of graph nodes and a point in the grid of points if the point lies within a hypercube of the rule corresponding to the graph node.
13 . The non-transitory computer readable storage medium of claim 8 wherein the graph structure is a bipartite graph comprising first and second sides, wherein the first side includes a plurality of graph nodes corresponding the plurality of rules, wherein the second side includes a plurality of range trees corresponding to the rule fields of the rule set, and wherein the bipartite graph includes an edge between a graph node in the plurality of graph nodes and a node in a range tree corresponding to a rule field R if a matching pattern specified for rule field R in the rule corresponding to the graph node falls within a value range associated with the node.
14 . The non-transitory computer readable storage medium of claim 8 wherein the graph structure comprises a plurality of range trees corresponding to the rule fields of the rule set, and wherein at least one node of at least one range tree includes a number of rules in the rule set that match a value range associated with the node.
15 . A computer system implementing a deep reinforcement learning (RL) system, comprising:
a processor; and
a non-transitory computer readable medium having stored thereon program code that, when executed, causes the processor to:
train an agent to construct an efficient decision tree for classifying network packets according to a rule set comprising a plurality of rules, the training including:
identifying, by an environment of the deep RL system, a leaf node in a decision tree;
computing, by the environment, a graph structure representing a state of the leaf node, wherein the graph structure includes information regarding how one or more rules in the rule set that are contained in the leaf node are distributed in a hypercube of the leaf node, wherein the hypercube is a convex shape in a multi-dimensional space having dimensions corresponding to rule fields in the rule set, and wherein the hypercube's boundaries are defined by split conditions used to reach the leaf node in the decision tree;
communicating, by the environment, the graph structure to the agent;
providing, by the agent, the graph structure as input to a graph neural network, the graph neural network being designed to accept variable-sized graph structures as input; and
generating, by the graph neural network based on the graph structure, an action to be taken on the leaf node for extending the decision tree.
16 . The computer system of claim 15 wherein the information regarding how the one or more rules contained in the leaf node are distributed in the hypercube of the leaf node comprises, for each of the one or more rules, information regarding how a hypercube of said each rule is placed in the hypercube of the leaf node, the hypercube of said each rule being a convex shape in the multi-dimensional space having boundaries defined by matching patterns included in said each rule.
17 . The computer system of claim 15 wherein providing, by the agent, the graph structure as input to a graph neural network comprises:
transforming, by the agent, the graph structure using a graph transformation function; and
providing, by the agent, the transformed graph structure as input to the graph neural network.
18 . The computer system of claim 15 wherein the action splits the leaf node into a plurality of child nodes in accordance with one or more split conditions defined along one of the rule fields.
19 . The computer system of claim 15 wherein the graph structure is a bipartite graph comprising first and second sides, wherein the first side includes a plurality of graph nodes corresponding the plurality of rules, wherein the second side includes a grid of points in the hypercube of the leaf node, and wherein the bipartite graph includes an edge between a graph node in the plurality of graph nodes and a point in the grid of points if the point lies within a hypercube of the rule corresponding to the graph node.
20 . The computer system of claim 15 wherein the graph structure is a bipartite graph comprising first and second sides, wherein the first side includes a plurality of graph nodes corresponding the plurality of rules, wherein the second side includes a plurality of range trees corresponding to the rule fields of the rule set, and wherein the bipartite graph includes an edge between a graph node in the plurality of graph nodes and a node in a range tree corresponding to a rule field R if a matching pattern specified for rule field R in the rule corresponding to the graph node falls within a value range associated with the node.
21 . The computer system of claim 15 wherein the graph structure comprises a plurality of range trees corresponding to the rule fields of the rule set, and wherein at least one node of at least one range tree includes a number of rules in the rule set that match a value range associated with the node.