Switch from theta to total energy as state variable

[trontrytel]

I think it would be good to use the same state variables as amip.

[trontrytel]

Or better yet, have both as options for testing