app/scripts/latex-importer/input/_minted/62B8750C0ACEBDA39A95140434E540A8.highlight.minted

\begin{MintedVerbatim}[commandchars=\\\{\}]
\PYG{k+kn}{import}\PYG{+w}{ }\PYG{n+nn}{torch}
\PYG{k+kn}{from}\PYG{+w}{ }\PYG{n+nn}{lerobot}\PYG{n+nn}{.}\PYG{n+nn}{datasets}\PYG{n+nn}{.}\PYG{n+nn}{lerobot\PYGZus{}dataset}\PYG{+w}{ }\PYG{k+kn}{import} \PYG{n}{LeRobotDataset}
\PYG{k+kn}{from}\PYG{+w}{ }\PYG{n+nn}{lerobot}\PYG{n+nn}{.}\PYG{n+nn}{datasets}\PYG{n+nn}{.}\PYG{n+nn}{streaming\PYGZus{}dataset}\PYG{+w}{ }\PYG{k+kn}{import} \PYG{n}{StreamingLeRobotDataset}

\PYG{n}{delta\PYGZus{}timestamps} \PYG{o}{=} \PYG{p}{\PYGZob{}}
    \PYG{l+s+s2}{\PYGZdq{}}\PYG{l+s+s2}{observation.images.wrist\PYGZus{}camera}\PYG{l+s+s2}{\PYGZdq{}}\PYG{p}{:} \PYG{p}{[}\PYG{o}{\PYGZhy{}}\PYG{l+m+mf}{0.2}\PYG{p}{,} \PYG{o}{\PYGZhy{}}\PYG{l+m+mf}{0.1}\PYG{p}{,} \PYG{l+m+mf}{0.0}\PYG{p}{]}  \PYG{c+c1}{\PYGZsh{} 0.2, and 0.1 seconds *before* each frame}
\PYG{p}{\PYGZcb{}}

\PYG{c+c1}{\PYGZsh{} Optionally, use StreamingLeRobotDataset to avoid downloading the dataset}
\PYG{n}{dataset} \PYG{o}{=} \PYG{n}{LeRobotDataset}\PYG{p}{(}
    \PYG{l+s+s2}{\PYGZdq{}}\PYG{l+s+s2}{lerobot/svla\PYGZus{}so101\PYGZus{}pickplace}\PYG{l+s+s2}{\PYGZdq{}}\PYG{p}{,}
    \PYG{n}{delta\PYGZus{}timestamps}\PYG{o}{=}\PYG{n}{delta\PYGZus{}timestamps}
\PYG{p}{)}

\PYG{c+c1}{\PYGZsh{} Streams frames from the Hugging Face Hub without loading into memory}
\PYG{n}{streaming\PYGZus{}dataset} \PYG{o}{=} \PYG{n}{StreamingLeRobotDataset}\PYG{p}{(}
    \PYG{l+s+s2}{\PYGZdq{}}\PYG{l+s+s2}{lerobot/svla\PYGZus{}so101\PYGZus{}pickplace}\PYG{l+s+s2}{\PYGZdq{}}\PYG{p}{,}
    \PYG{n}{delta\PYGZus{}timestamps}\PYG{o}{=}\PYG{n}{delta\PYGZus{}timestamps}
\PYG{p}{)}

\PYG{c+c1}{\PYGZsh{} Get the 100th frame in the dataset by }
\PYG{n}{sample} \PYG{o}{=} \PYG{n}{dataset}\PYG{p}{[}\PYG{l+m+mi}{100}\PYG{p}{]}
\PYG{n+nb}{print}\PYG{p}{(}\PYG{n}{sample}\PYG{p}{)}
\PYG{c+c1}{\PYGZsh{} \PYGZob{}}
\PYG{c+c1}{\PYGZsh{} \PYGZsq{}observation.state\PYGZsq{}: tensor([...]), }
\PYG{c+c1}{\PYGZsh{} \PYGZsq{}action\PYGZsq{}: tensor([...]), }
\PYG{c+c1}{\PYGZsh{} \PYGZsq{}observation.images.wrist\PYGZus{}camera\PYGZsq{}: tensor([3, C, H, W]), for delta timesteps}
\PYG{c+c1}{\PYGZsh{} ...}
\PYG{c+c1}{\PYGZsh{} \PYGZcb{}}

\PYG{n}{batch\PYGZus{}size}\PYG{o}{=}\PYG{l+m+mi}{16}
\PYG{c+c1}{\PYGZsh{} wrap the dataset in a DataLoader to use process it batches for training purposes}
\PYG{n}{data\PYGZus{}loader} \PYG{o}{=} \PYG{n}{torch}\PYG{o}{.}\PYG{n}{utils}\PYG{o}{.}\PYG{n}{data}\PYG{o}{.}\PYG{n}{DataLoader}\PYG{p}{(}
    \PYG{n}{dataset}\PYG{p}{,}
    \PYG{n}{batch\PYGZus{}size}\PYG{o}{=}\PYG{n}{batch\PYGZus{}size}
\PYG{p}{)}

\PYG{c+c1}{\PYGZsh{} Iterate over the DataLoader in a training loop}
\PYG{n}{num\PYGZus{}epochs} \PYG{o}{=} \PYG{l+m+mi}{1}
\PYG{n}{device} \PYG{o}{=} \PYG{l+s+s2}{\PYGZdq{}}\PYG{l+s+s2}{cuda}\PYG{l+s+s2}{\PYGZdq{}} \PYG{k}{if} \PYG{n}{torch}\PYG{o}{.}\PYG{n}{cuda}\PYG{o}{.}\PYG{n}{is\PYGZus{}available}\PYG{p}{(}\PYG{p}{)} \PYG{k}{else} \PYG{l+s+s2}{\PYGZdq{}}\PYG{l+s+s2}{cpu}\PYG{l+s+s2}{\PYGZdq{}}

\PYG{k}{for} \PYG{n}{epoch} \PYG{o+ow}{in} \PYG{n+nb}{range}\PYG{p}{(}\PYG{n}{num\PYGZus{}epochs}\PYG{p}{)}\PYG{p}{:}
    \PYG{k}{for} \PYG{n}{batch} \PYG{o+ow}{in} \PYG{n}{data\PYGZus{}loader}\PYG{p}{:}
        \PYG{c+c1}{\PYGZsh{} Move data to the appropriate device (e.g., GPU)}
        \PYG{n}{observations} \PYG{o}{=} \PYG{n}{batch}\PYG{p}{[}\PYG{l+s+s2}{\PYGZdq{}}\PYG{l+s+s2}{observation.state}\PYG{l+s+s2}{\PYGZdq{}}\PYG{p}{]}\PYG{o}{.}\PYG{n}{to}\PYG{p}{(}\PYG{n}{device}\PYG{p}{)}
        \PYG{n}{actions} \PYG{o}{=} \PYG{n}{batch}\PYG{p}{[}\PYG{l+s+s2}{\PYGZdq{}}\PYG{l+s+s2}{action}\PYG{l+s+s2}{\PYGZdq{}}\PYG{p}{]}\PYG{o}{.}\PYG{n}{to}\PYG{p}{(}\PYG{n}{device}\PYG{p}{)}
        \PYG{n}{images} \PYG{o}{=} \PYG{n}{batch}\PYG{p}{[}\PYG{l+s+s2}{\PYGZdq{}}\PYG{l+s+s2}{observation.images.wrist\PYGZus{}camera}\PYG{l+s+s2}{\PYGZdq{}}\PYG{p}{]}\PYG{o}{.}\PYG{n}{to}\PYG{p}{(}\PYG{n}{device}\PYG{p}{)}

        \PYG{c+c1}{\PYGZsh{} Next, you can do amazing\PYGZus{}model.forward(batch)}
        \PYG{o}{.}\PYG{o}{.}\PYG{o}{.}
\end{MintedVerbatim}