Unifying Robot Optimization: Monte Carlo Tree Search with Tensor Factorization

Teng Xue, Amirreza Razmjoo†, Yan Zhang†, Sylvain Calinon

Idiap Research Institute and EPFL

Paper

Code

Abstract

Many robotic tasks, such as inverse kinematics, motion planning, and optimal control, can be formulated as optimization problems. Solving these problems involves addressing nonlinear kinematics, complex contact dynamics, and long-horizon planning, each posing distinct challenges for state-of-the-art optimization methods. To efficiently solve a wide range of tasks across varying scenarios, researchers either develop specialized algorithms for the task to achieve, or switch between different frameworks. Monte Carlo Tree Search (MCTS) is a general-purpose decision-making tool that enables strategic exploration across problem instances without relying on task-specific structures. However, MCTS suffers from combinatorial complexity, leading to slow convergence and high memory usage. To address this limitation, we propose Tensor Train Tree Search (TTTS), which leverages tensor factorization to exploit the separable structure of decision trees. This yields a low-rank, linear-complexity representation that significantly reduces both computation time and storage requirements. We prove that TTTS can efficiently reach the bounded global optimum within a finite time. Experimental results across inverse kinematics, motion planning, multi-stage motion planning around obstacles, and bimanual whole-body manipulation demonstrate the efficiency of TTTS to discover diverse sets of optimal solutions.

Method

We consider a general robot optimization formulation that can address diverse problems, such as inverse kinematics, motion planning, multi-stage motion planning with mode switching, and model predictive control, which corresponds to a large sets of mathematical programs, including nonlinear programming (NLP), large-scale (aka. high-dim) NLP and mixed-integer nonlinear programming (MINLP). In particular, we leverage basis functions to reduce the dimensionality for large-scale NLP formulation.

The unified mathematical formulation for our robot optimization problems is:

where,

(2) encodes the full trajectory of decision variables using basis functions.

(3) enforces allowed transitions in the discrete mode set given the current mode and the system trajectory.

(4) and (5) represent the system dynamics, consistency of different modes, and other physical constraints.

We propose TTTS to tackle the above MINLP problem. Tensor Train (TT) is introduced as an efficient representation for decision trees that significantly reduces combinatorial complexity to linear complexity. A multi-layer decision tree can be represented as a high-dimensional tensor, as shown in Figure 1(A). Since tensor factorization can be used to represent a high-dim tensor through a chain of third-order tensor cores (Figure 1(B)), resulting in TT format, a multi-layer node-based tree can be expressed as TT-Tree, which significantly reduces the combinatorial complexity to linear complexity in both computation time and storage space. With such structure, the node value and visiting info can efficiently computed, as shown in Figure 1(C), which illustrates an example about how how to compute the value of a node in the first layer of a three-layer tree using TT cores.

Figure 1: Tree-Tensor-TT transformation.