Use WMStateHelper to waitFor transition

Using the UI automator to wait for elements is unstable and occasionally triggers `android.os.DeadObjectException`

Using the activity manager state to wait for the device is a more reliable way of doing this

Bug: 175877007
Test: atest WMShellFlickerTests:com.android.wm.shell.flicker.pip.PipKeyboardTest -- --abi arm64-v8a
Change-Id: I0c838ea91e18cb6e9821efd85b05f507b131cc71