Add preliminary docstrings to viz.py

stac_mjx/viz.py

@@ -1,3 +1,7 @@
+"""A collection mujoco-mjx vizualization utilities.
+"""
+
+
 from dm_control import mjcf
 from dm_control.locomotion.walkers import rescale
 from dm_control.mujoco.wrapper.mjbindings import enums
@@ -192,6 +196,8 @@ def overlay_viz(
     save_path,
     camera: Text = "close_profile",
 ):
+    """TBD
+    """
     scene_option = wrapper.MjvOption()
     # scene_option.geomgroup[1] = 0
     scene_option.geomgroup[2] = 1
@@ -302,6 +308,8 @@ def overlay_viz(
 
 
 def mujoco_viz(data_path, model_xml, n_frames, save_path, start_frame: int = 0):
+    """Renders forward kinematics from keypoint positions.
+    """
     scene_option = mujoco.MjvOption()
     # scene_option.geomgroup[1] = 0
     scene_option.geomgroup[2] = 1