From 614ea56208e28b919db8f444c90599d4c51dc89e Mon Sep 17 00:00:00 2001
From: Jason Foat <jfoat@salk.edu>
Date: Mon, 22 Jul 2024 18:42:30 +0000
Subject: [PATCH] Add preliminary docstrings to viz.py

---
 stac_mjx/viz.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/stac_mjx/viz.py b/stac_mjx/viz.py
index 9f2c1e8..262b842 100644
--- a/stac_mjx/viz.py
+++ b/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