Fix M:initial_state/0 handling

`initial_state/0` may return `{'foo_state', Data}` or `'foo_state'`.  Adjust
the caller accordingly.