Update README.md

README.md

@@ -19,13 +19,25 @@ We envision our model as a platform for fruit fly biophysics simulations and for
 
 The fruit fly body model lives in [this directory][fly-home]. To visualize it, you can drag-and-drop `fruitfly.xml` or `floor.xml` to MuJoCo's `simulate` viewer.
 
-Beginning interacting with the model via Python is as simple as:
+Interacting with the fly via Python is as simple as:
 
 ```python
-from dm_control import mujoco
+import numpy as np
+import mediapy
 
-physics = mujoco.Physics.from_xml_path('flybody/fruitfly/assets/fruitfly.xml')  # Load model.
-physics.step()  # Step simulation.
+from flybody.fly_envs import walk_imitation
+
+# Create walking imitation environment.
+env = walk_imitation()
+
+# Run environment loop with random actions for a bit.
+for _ in range(100):
+   action = np.random.normal(size=59)  # 59 is the walking action dimension.
+   timestep = env.step(action)
+
+# Generate a pretty image.
+pixels = env.physics.render(camera_id=1)
+mediapy.show_image(pixels)
 ```
 
 The quickest way to get started with `flybody` is to take a look at a [tutorial notebook][tutorial] or [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)][tutorial-colab].