There are some misunderstandings on how the UIDocument component works here, which is due to our lack of documentation, so apologies about that.
The UXML file won't be loaded by the UIDocument component until OnEnable runs on it, so trying to access the root during Start in the NewBehaviourScript component means you may not get the root of the cloned tree. We highly suggest using OnEnable on your component to guarantee the expected behavior.

In any case, when you wait for the end of the frame through your yield return statement, that means one frame will be drawn and that one frame can contain the cloned tree. Also if I understand correctly, you want to use a USS transition to move the square after the first frame, and the preferred way of moving something on screen like that would be to register a callback to the GeometryChangedEvent and change the position of the square in the callback. That guarantees the transition will kick in since the first layout will have the square in place and when the new position is set, it'll transition to it as expected.