I agree with Gary. Time is not sufficient to travel from Budapest to Vienna and back in one day and at the same time visit the beautiful capital city of Austria.
I suggest you do a one day trip into the Hungarian Puszta or eventually to Lake Balaton. The Hungarian Railway Company MAV will certainly be able to help you in that
www.mav.hu.
A trip on the Danube with a boat may also be a great thing to do.
Enjoy your trip.
Ramón