I’m not sure if you still need an answer, but I just came across the same issue myself. The way I got it working was by subclassing AjaxTreeModel and overriding setRootTreeNode. All the expanded/collapsed nodes get reset when a different rootTreeNode is set.
@Override
public void setRootTreeNode(Object rootTreeNode) {
if (rootTreeNode != rootTreeNode()) {
super.setRootTreeNode(rootTreeNode);
// expand tree nodes here
}
}