Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

docs/source/recorders.rst

@@ -0,0 +1,121 @@
+.. _recorders:
+
+Recorders
+=========
+
+The recorder manager provides lifecycle hooks for logging data during
+rollouts. Unlike rewards, recorders have no effect on the optimization
+loop. They exist purely to let you capture observations, actions, or any
+other environment state without modifying mjlab internals.
+
+Each recorder term is a class that you implement. mjlab calls its methods
+at the right moments and leaves all I/O decisions to you. If the ``recorders``
+dictionary on ``ManagerBasedRlEnvCfg`` is empty, the environment substitutes a
+lightweight no-op manager with zero overhead.
+
+
+Lifecycle hooks
+---------------
+
+The manager exposes three hooks per term:
+
+``record_pre_reset(env_ids)``
+    Called inside ``env.step()`` before terminated environments are reset.
+    ``obs_buf`` holds the observation from the end of the *previous* step
+    (the input the agent used to choose the terminal action, not the
+    post-action terminal state). ``action_manager.action`` holds the
+    terminal action and is still valid here; it will be zeroed for these
+    environments by ``_reset_idx`` immediately after. ``reward_buf`` holds
+    the terminal reward. This is the right place to record the terminal
+    transition ``(obs_t, action_t, reward_t, done=True)``.
+
+``record_post_reset(env_ids)``
+    Called after a reset completes and fresh observations are available.
+    Fires at the end of ``env.reset()`` (all environments) and within
+    ``env.step()`` after each batch of terminated environments is reset.
+    ``obs_buf[env_ids]`` holds the initial observation of the new episode;
+    ``action_manager.action[env_ids]`` is zero. Use this to initialize
+    per-episode state or record the first observation.
+
+``record_post_step()``
+    Called at the end of every ``env.step()`` with fresh observations.
+    For environments that reset during this step, ``action_manager.action``
+    has been zeroed and ``obs_buf`` holds the initial state of the new
+    episode rather than the post-action terminal observation. Use
+    ``record_pre_reset`` for those environments' terminal transitions and
+    ``self._env.reset_buf`` to identify which environments reset.
+
+``close()``
+    Called when the environment closes. Release file handles and flush
+    buffers here.
+
+
+Writing a recorder term
+------------------------
+
+Subclass :class:`~mjlab.managers.RecorderTerm` and override whichever
+hooks you need. The environment is available as ``self._env``, giving
+access to ``self._env.obs_buf``, ``self._env.action_manager.action``, and
+all other managers.
+
+.. code-block:: python
+
+    import csv
+    from mjlab.managers import RecorderTerm, RecorderTermCfg
+
+    class CsvRecorder(RecorderTerm):
+        def __init__(self, cfg, env):
+            super().__init__(cfg, env)
+            self._file = open(cfg.params["path"], "w", newline="")
+            self._writer = csv.writer(self._file)
+
+        def record_pre_reset(self, env_ids):
+            # Terminal transition: action is still intact here.
+            # It will be zeroed by _reset_idx immediately after this returns.
+            obs = self._env.obs_buf["actor"][env_ids].cpu().numpy()
+            act = self._env.action_manager.action[env_ids].cpu().numpy()
+            for o, a in zip(obs, act):
+                self._writer.writerow(o.tolist() + a.tolist())
+
+        def record_post_step(self):
+            # Skip envs that just reset: their terminal pair was written
+            # in record_pre_reset and their action is now zeroed.
+            mask = ~self._env.reset_buf
+            obs = self._env.obs_buf["actor"][mask].cpu().numpy()
+            act = self._env.action_manager.action[mask].cpu().numpy()
+            for o, a in zip(obs, act):
+                self._writer.writerow(o.tolist() + a.tolist())
+
+        def close(self):
+            self._file.close()
+
+The term receives the full ``cfg`` object so it can read any values you
+put in ``cfg.params``.
+
+
+Registration
+------------
+
+Add the term to the ``recorders`` dictionary on your environment config:
+
+.. code-block:: python
+
+    from dataclasses import dataclass, field
+    from mjlab.managers import RecorderTermCfg
+
+    @dataclass
+    class MyEnvCfg(SomeTaskEnvCfg):
+        recorders: dict = field(default_factory=lambda: {
+            "csv": RecorderTermCfg(
+                func=CsvRecorder,
+                params={"path": "rollout.csv"},
+            )
+        })
+
+Multiple terms can be registered under different keys and run together.
+
+.. note::
+
+    ``func`` must be a :class:`~mjlab.managers.RecorderTerm` subclass.
+    Function-based terms are not supported because recorder terms are
+    stateful.