How can I improve this RL algorithm? This is the task: I have a unitary target matrix $T$, to decompose using matrices from a fixed (finite) universal set $\{M_i\}$, e.g. $T = M_6M_3M_9M_0$. The set is universal in the sense that one can reach any unitary with by combining its elements appropriately.

I wrote an RL algorithm with a policy network that I feed with the target $T$ and with the current "state" (i.e. the current product of the matrices that the policy chose), which returns the index of which matrix to pick next. In pseudocode, this is what I wrote:

reward = 0
 state = Identity_matrix
 chosen_indices = []
 for(i=0,max_episode_length,i++):
    prob_dist = policy(state,T)
    k = random_int(prob_dist)
    state = state * M[k]
    reward -= 1
    if ||T-state|| == 0:
        reward += 100
        end_episode()
loss = discounted_rewards * cross_entropy([prob_dist],[one_hot(k)])
minimize(loss)


I repeat this loop for several random examples of $T$ of which I know the decomposition. But, even if I limit $T$ to be a single matrix from the set, the policy learns a couple of them and then it stops learning, i.e. when I pick a new $T$, the average episode reward remains around zero.

I tried fiddling with the learning rate, the discount factor, the width of the policy network, nothing seems to make things work. What am I doing wrong?
There's no reason to think that reinforcement learning would be effective here.  Indeed, the surprise would be if it did work!  Machine learning isn't a magic silver bullet.  I wouldn't expect it to be of any assistance here; the problem relates to deep questions about the structure of the general linear group, which might be quite challenging for your policy network to learn.

Your problem is known as the membership problem for the general linear group, which might be hard (i.e., no efficient algorithm is known), depending on which field you are working over.  See https://cstheory.stackexchange.com/q/38000/5038.  If we take $T=0$ and are working with $n \times n$ matrices, this is called the matrix mortality problem; it is undecidable for $n \ge 3$.  See also https://en.wikipedia.org/wiki/List_of_undecidable_problems#Problems_about_matrices.  So, in the general case, your problem is undecidable.

Obviously, you shouldn't expect your approach to successfully solve an undecidable problem.