I've written this small generic library for the purpose of my Bachelor's thesis. It's fully functional and unit tested and I want to get as many opinions as possible regarding overall code quality (readability, naming, structure, coupling, cohesion, etc.).
GitHub (this project is licensed under the terms of the MIT license)
Usage
Mcts<State, Action, Player> mcts = Mcts.initializeIterations(NUMBER_OF_ITERATIONS);
Action mostPromisingAction = mcts.uctSearchWithExploration(state, explorationParameter);
Interfaces
public interface MctsDomainState<ActionT, AgentT extends MctsDomainAgent> {
boolean isTerminal();
AgentT getCurrentAgent();
AgentT getPreviousAgent();
int getNumberOfAvailableActionsForCurrentAgent();
List<ActionT> getAvailableActionsForCurrentAgent();
MctsDomainState performActionForCurrentAgent(ActionT action);
MctsDomainState skipCurrentAgent();
}
public interface MctsDomainAgent<StateT extends MctsDomainState> {
StateT getTerminalStateByPerformingSimulationFromState(StateT state);
double getRewardFromTerminalState(StateT terminalState);
}
Class that iteratively builds the search tree
public class Mcts<StateT extends MctsDomainState<ActionT, AgentT>, ActionT, AgentT extends MctsDomainAgent<StateT>> {
private static final double NO_EXPLORATION = 0;
private final int numberOfIterations;
private double explorationParameter;
private final Cloner cloner;
public static<StateT extends MctsDomainState<ActionT, AgentT>, ActionT, AgentT extends MctsDomainAgent<StateT>>
Mcts<StateT, ActionT, AgentT> initializeIterations(int numberOfIterations) {
Cloner cloner = new Cloner();
return new Mcts<>(numberOfIterations, cloner);
}
private Mcts(int numberOfIterations, Cloner cloner) {
this.numberOfIterations = numberOfIterations;
this.cloner = cloner;
}
public void dontClone(final Class<?>... classes) {
cloner.dontClone(classes);
}
public ActionT uctSearchWithExploration(StateT state, double explorationParameter) {
setExplorationForSearch(explorationParameter);
MctsTreeNode<StateT, ActionT, AgentT> rootNode = new MctsTreeNode<>(state, cloner);
for (int i = 0; i < numberOfIterations; i++) {
performMctsIteration(rootNode, state.getCurrentAgent());
}
return getNodesMostPromisingAction(rootNode);
}
private void setExplorationForSearch(double explorationParameter) {
this.explorationParameter = explorationParameter;
}
private void performMctsIteration(MctsTreeNode<StateT, ActionT, AgentT> rootNode, AgentT agentInvoking) {
MctsTreeNode<StateT, ActionT, AgentT> selectedChildNode = treePolicy(rootNode);
StateT terminalState = getTerminalStateFromDefaultPolicy(selectedChildNode, agentInvoking);
backPropagate(selectedChildNode, terminalState);
}
private MctsTreeNode<StateT, ActionT, AgentT> treePolicy(MctsTreeNode<StateT, ActionT, AgentT> node) {
while (!node.representsTerminalState()) {
if (!node.representedStatesCurrentAgentHasAvailableActions())
return expandWithoutAction(node);
else if (!node.isFullyExpanded())
return expandWithAction(node);
else
node = getNodesBestChild(node);
}
return node;
}
private MctsTreeNode<StateT, ActionT, AgentT> expandWithoutAction(MctsTreeNode<StateT, ActionT, AgentT> node) {
return node.addNewChildWithoutAction();
}
private MctsTreeNode<StateT, ActionT, AgentT> expandWithAction(MctsTreeNode<StateT, ActionT, AgentT> node) {
ActionT randomUntriedAction = getRandomActionFromNodesUntriedActions(node);
return node.addNewChildFromAction(randomUntriedAction);
}
private ActionT getRandomActionFromNodesUntriedActions(MctsTreeNode<StateT, ActionT, AgentT> node) {
List<ActionT> untriedActions = node.getUntriedActionsForCurrentAgent();
Collections.shuffle(untriedActions);
return untriedActions.get(0);
}
private MctsTreeNode<StateT, ActionT, AgentT> getNodesBestChild(MctsTreeNode<StateT, ActionT, AgentT> node) {
validateBestChildComputable(node);
return getNodesBestChildConfidentlyWithExploration(node, explorationParameter);
}
private void validateBestChildComputable(MctsTreeNode<StateT, ActionT, AgentT> node) {
if (!node.hasChildNodes())
throw new UnsupportedOperationException("Error: operation not supported if child nodes empty");
else if (!node.isFullyExpanded())
throw new UnsupportedOperationException("Error: operation not supported if node not fully expanded");
else if (node.hasUnvisitedChild())
throw new UnsupportedOperationException(
"Error: operation not supported if node contains an unvisited child");
}
private ActionT getNodesMostPromisingAction(MctsTreeNode<StateT, ActionT, AgentT> node) {
validateBestChildComputable(node);
MctsTreeNode<StateT, ActionT, AgentT> bestChildWithoutExploration =
getNodesBestChildConfidentlyWithExploration(node, NO_EXPLORATION);
return bestChildWithoutExploration.getIncomingAction();
}
private MctsTreeNode<StateT, ActionT, AgentT> getNodesBestChildConfidentlyWithExploration(
MctsTreeNode<StateT, ActionT, AgentT> node, double explorationParameter) {
return node.getChildNodes().stream()
.max((node1, node2) -> Double.compare(
calculateUctValue(node1, explorationParameter),
calculateUctValue(node2, explorationParameter))).get();
}
private double calculateUctValue(MctsTreeNode<StateT, ActionT, AgentT> node, double explorationParameter) {
return node.getDomainTheoreticValue()
+ explorationParameter
* (Math.sqrt((2 * Math.log(node.getParentsVisitCount())) / node.getVisitCount()));
}
private StateT getTerminalStateFromDefaultPolicy(
MctsTreeNode<StateT, ActionT, AgentT> node, AgentT agentInvoking) {
StateT nodesStateClone = node.getDeepCloneOfRepresentedState();
return agentInvoking.getTerminalStateByPerformingSimulationFromState(nodesStateClone);
}
private void backPropagate(MctsTreeNode<StateT, ActionT, AgentT> node, StateT terminalState) {
while (node != null) {
updateNodesDomainTheoreticValue(node, terminalState);
node = node.getParentNode();
}
}
private void updateNodesDomainTheoreticValue(MctsTreeNode<StateT, ActionT, AgentT> node, StateT terminalState) {
// violation of the law of demeter
AgentT parentsStatesCurrentAgent = node.getRepresentedStatesPreviousAgent();
double reward = parentsStatesCurrentAgent.getRewardFromTerminalState(terminalState);
node.updateDomainTheoreticValue(reward);
}
}
Class that represents the tree node
class MctsTreeNode<StateT extends MctsDomainState<ActionT, AgentT>, ActionT, AgentT extends MctsDomainAgent> {
private final MctsTreeNode<StateT, ActionT, AgentT> parentNode;
private final ActionT incomingAction;
private final StateT representedState;
private int visitCount;
private double totalReward;
private List<MctsTreeNode<StateT, ActionT, AgentT>> childNodes;
private final Cloner cloner;
protected MctsTreeNode(StateT representedState, Cloner cloner) {
this(null, null, representedState, cloner);
}
private MctsTreeNode(MctsTreeNode<StateT, ActionT, AgentT> parentNode, ActionT incomingAction,
StateT representedState, Cloner cloner) {
this.parentNode = parentNode;
this.incomingAction = incomingAction;
this.representedState = representedState;
this.visitCount = 0;
this.totalReward = 0.0;
this.childNodes = new ArrayList<>();
this.cloner = cloner;
}
protected MctsTreeNode<StateT, ActionT, AgentT> getParentNode() {
return parentNode;
}
protected ActionT getIncomingAction() {
return incomingAction;
}
protected int getVisitCount() {
return visitCount;
}
protected int getParentsVisitCount() {
return parentNode.getVisitCount();
}
protected List<MctsTreeNode<StateT, ActionT, AgentT>> getChildNodes() {
return childNodes;
}
protected boolean hasChildNodes() {
return childNodes.size() > 0;
}
protected boolean representsTerminalState() {
return representedState.isTerminal();
}
protected AgentT getRepresentedStatesPreviousAgent() {
return representedState.getPreviousAgent();
}
protected boolean representedStatesCurrentAgentHasAvailableActions() {
return representedState.getNumberOfAvailableActionsForCurrentAgent() > 0;
}
protected boolean isFullyExpanded() {
return representedState.getNumberOfAvailableActionsForCurrentAgent() == childNodes.size();
}
protected boolean hasUnvisitedChild () {
return childNodes.stream()
.anyMatch(MctsTreeNode::isUnvisited);
}
private boolean isUnvisited() {
return visitCount == 0;
}
protected MctsTreeNode<StateT, ActionT, AgentT> addNewChildWithoutAction() {
StateT childNodeState = getDeepCloneOfRepresentedState();
childNodeState.skipCurrentAgent();
return appendNewChildInstance(childNodeState, null);
}
protected MctsTreeNode<StateT, ActionT, AgentT> addNewChildFromAction(ActionT action) {
if (!isUntriedAction(action))
throw new IllegalArgumentException("Error: invalid action passed as function parameter");
else
return addNewChildFromUntriedAction(action);
}
private boolean isUntriedAction(ActionT action) {
return getUntriedActionsForCurrentAgent().contains(action);
}
protected List<ActionT> getUntriedActionsForCurrentAgent() {
List<ActionT> availableActions = representedState.getAvailableActionsForCurrentAgent();
List<ActionT> untriedActions = new ArrayList<>(availableActions);
List<ActionT> triedActions = getTriedActionsForCurrentAgent();
untriedActions.removeAll(triedActions);
return untriedActions;
}
private List<ActionT> getTriedActionsForCurrentAgent() {
return childNodes.stream()
.map(MctsTreeNode::getIncomingAction)
.collect(Collectors.toList());
}
private MctsTreeNode<StateT, ActionT, AgentT> addNewChildFromUntriedAction(ActionT incomingAction) {
StateT childNodeState = getNewStateFromAction(incomingAction);
return appendNewChildInstance(childNodeState, incomingAction);
}
private StateT getNewStateFromAction(ActionT action) {
StateT representedStateClone = getDeepCloneOfRepresentedState();
representedStateClone.performActionForCurrentAgent(action);
return representedStateClone;
}
protected StateT getDeepCloneOfRepresentedState() {
return cloner.deepClone(representedState);
}
private MctsTreeNode<StateT, ActionT, AgentT> appendNewChildInstance(
StateT representedState, ActionT incomingAction) {
MctsTreeNode<StateT, ActionT, AgentT> childNode = new MctsTreeNode<>(
this, incomingAction, representedState, cloner);
childNodes.add(childNode);
return childNode;
}
protected void updateDomainTheoreticValue(double rewardAddend) {
visitCount += 1;
totalReward += rewardAddend;
}
protected double getDomainTheoreticValue() {
return totalReward / visitCount;
}
}