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?