Initial commit

.gitattributes

@@ -33,3 +33,19 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/adept_models/kitchen/textures/marble1.png filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/adept_models/kitchen/textures/metal1.png filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/adept_models/kitchen/textures/tile1.png filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/adept_models/kitchen/textures/wood1.png filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/adept_models/scenes/textures/white_marble_tile.png filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/adept_models/scenes/textures/white_marble_tile2.png filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/third_party/franka/franka_panda.png filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/hand.stl filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link0.stl filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link1.stl filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link2.stl filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link3.stl filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link4.stl filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link5.stl filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link6.stl filter=lfs diff=lfs merge=lfs -text
+env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link7.stl filter=lfs diff=lfs merge=lfs -text

codecs/imagecodecs_numcodecs.py

@@ -0,0 +1,1386 @@
+
+# imagecodecs/numcodecs.py
+
+# Copyright (c) 2021-2022, Christoph Gohlke
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice,
+#    this list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+#    contributors may be used to endorse or promote products derived from
+#    this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+
+"""Additional numcodecs implemented using imagecodecs."""
+
+__version__ = '2022.9.26'
+
+__all__ = ('register_codecs',)
+
+import numpy
+from numcodecs.abc import Codec
+from numcodecs.registry import register_codec, get_codec
+
+import imagecodecs
+
+
+def protective_squeeze(x: numpy.ndarray):
+    """
+    Squeeze dim only if it's not the last dim.
+    Image dim expected to be *, H, W, C
+    """
+    img_shape = x.shape[-3:]
+    if len(x.shape) > 3:
+        n_imgs = numpy.prod(x.shape[:-3])
+        if n_imgs > 1:
+            img_shape = (-1,) + img_shape
+    return x.reshape(img_shape)
+
+def get_default_image_compressor(**kwargs):
+    if imagecodecs.JPEGXL:
+        # has JPEGXL
+        this_kwargs = {
+            'effort': 3,
+            'distance': 0.3,
+            # bug in libjxl, invalid codestream for non-lossless
+            # when decoding speed > 1
+            'decodingspeed': 1
+        }
+        this_kwargs.update(kwargs)
+        return JpegXl(**this_kwargs)
+    else:
+        this_kwargs = {
+            'level': 50
+        }
+        this_kwargs.update(kwargs)
+        return Jpeg2k(**this_kwargs)
+
+class Aec(Codec):
+    """AEC codec for numcodecs."""
+
+    codec_id = 'imagecodecs_aec'
+
+    def __init__(
+        self, bitspersample=None, flags=None, blocksize=None, rsi=None
+    ):
+        self.bitspersample = bitspersample
+        self.flags = flags
+        self.blocksize = blocksize
+        self.rsi = rsi
+
+    def encode(self, buf):
+        return imagecodecs.aec_encode(
+            buf,
+            bitspersample=self.bitspersample,
+            flags=self.flags,
+            blocksize=self.blocksize,
+            rsi=self.rsi,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.aec_decode(
+            buf,
+            bitspersample=self.bitspersample,
+            flags=self.flags,
+            blocksize=self.blocksize,
+            rsi=self.rsi,
+            out=_flat(out),
+        )
+
+
+class Apng(Codec):
+    """APNG codec for numcodecs."""
+
+    codec_id = 'imagecodecs_apng'
+
+    def __init__(self, level=None, photometric=None, delay=None):
+        self.level = level
+        self.photometric = photometric
+        self.delay = delay
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.apng_encode(
+            buf,
+            level=self.level,
+            photometric=self.photometric,
+            delay=self.delay,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.apng_decode(buf, out=out)
+
+
+class Avif(Codec):
+    """AVIF codec for numcodecs."""
+
+    codec_id = 'imagecodecs_avif'
+
+    def __init__(
+        self,
+        level=None,
+        speed=None,
+        tilelog2=None,
+        bitspersample=None,
+        pixelformat=None,
+        numthreads=None,
+        index=None,
+    ):
+        self.level = level
+        self.speed = speed
+        self.tilelog2 = tilelog2
+        self.bitspersample = bitspersample
+        self.pixelformat = pixelformat
+        self.numthreads = numthreads
+        self.index = index
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.avif_encode(
+            buf,
+            level=self.level,
+            speed=self.speed,
+            tilelog2=self.tilelog2,
+            bitspersample=self.bitspersample,
+            pixelformat=self.pixelformat,
+            numthreads=self.numthreads,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.avif_decode(
+            buf, index=self.index, numthreads=self.numthreads, out=out
+        )
+
+
+class Bitorder(Codec):
+    """Bitorder codec for numcodecs."""
+
+    codec_id = 'imagecodecs_bitorder'
+
+    def encode(self, buf):
+        return imagecodecs.bitorder_encode(buf)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.bitorder_decode(buf, out=_flat(out))
+
+
+class Bitshuffle(Codec):
+    """Bitshuffle codec for numcodecs."""
+
+    codec_id = 'imagecodecs_bitshuffle'
+
+    def __init__(self, itemsize=1, blocksize=0):
+        self.itemsize = itemsize
+        self.blocksize = blocksize
+
+    def encode(self, buf):
+        return imagecodecs.bitshuffle_encode(
+            buf, itemsize=self.itemsize, blocksize=self.blocksize
+        ).tobytes()
+
+    def decode(self, buf, out=None):
+        return imagecodecs.bitshuffle_decode(
+            buf,
+            itemsize=self.itemsize,
+            blocksize=self.blocksize,
+            out=_flat(out),
+        )
+
+
+class Blosc(Codec):
+    """Blosc codec for numcodecs."""
+
+    codec_id = 'imagecodecs_blosc'
+
+    def __init__(
+        self,
+        level=None,
+        compressor=None,
+        typesize=None,
+        blocksize=None,
+        shuffle=None,
+        numthreads=None,
+    ):
+        self.level = level
+        self.compressor = compressor
+        self.typesize = typesize
+        self.blocksize = blocksize
+        self.shuffle = shuffle
+        self.numthreads = numthreads
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.blosc_encode(
+            buf,
+            level=self.level,
+            compressor=self.compressor,
+            typesize=self.typesize,
+            blocksize=self.blocksize,
+            shuffle=self.shuffle,
+            numthreads=self.numthreads,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.blosc_decode(
+            buf, numthreads=self.numthreads, out=_flat(out)
+        )
+
+
+class Blosc2(Codec):
+    """Blosc2 codec for numcodecs."""
+
+    codec_id = 'imagecodecs_blosc2'
+
+    def __init__(
+        self,
+        level=None,
+        compressor=None,
+        typesize=None,
+        blocksize=None,
+        shuffle=None,
+        numthreads=None,
+    ):
+        self.level = level
+        self.compressor = compressor
+        self.typesize = typesize
+        self.blocksize = blocksize
+        self.shuffle = shuffle
+        self.numthreads = numthreads
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.blosc2_encode(
+            buf,
+            level=self.level,
+            compressor=self.compressor,
+            typesize=self.typesize,
+            blocksize=self.blocksize,
+            shuffle=self.shuffle,
+            numthreads=self.numthreads,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.blosc2_decode(
+            buf, numthreads=self.numthreads, out=_flat(out)
+        )
+
+
+class Brotli(Codec):
+    """Brotli codec for numcodecs."""
+
+    codec_id = 'imagecodecs_brotli'
+
+    def __init__(self, level=None, mode=None, lgwin=None):
+        self.level = level
+        self.mode = mode
+        self.lgwin = lgwin
+
+    def encode(self, buf):
+        return imagecodecs.brotli_encode(
+            buf, level=self.level, mode=self.mode, lgwin=self.lgwin
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.brotli_decode(buf, out=_flat(out))
+
+
+class ByteShuffle(Codec):
+    """ByteShuffle codec for numcodecs."""
+
+    codec_id = 'imagecodecs_byteshuffle'
+
+    def __init__(
+        self, shape, dtype, axis=-1, dist=1, delta=False, reorder=False
+    ):
+        self.shape = tuple(shape)
+        self.dtype = numpy.dtype(dtype).str
+        self.axis = axis
+        self.dist = dist
+        self.delta = bool(delta)
+        self.reorder = bool(reorder)
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        assert buf.shape == self.shape
+        assert buf.dtype == self.dtype
+        return imagecodecs.byteshuffle_encode(
+            buf,
+            axis=self.axis,
+            dist=self.dist,
+            delta=self.delta,
+            reorder=self.reorder,
+        ).tobytes()
+
+    def decode(self, buf, out=None):
+        if not isinstance(buf, numpy.ndarray):
+            buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape)
+        return imagecodecs.byteshuffle_decode(
+            buf,
+            axis=self.axis,
+            dist=self.dist,
+            delta=self.delta,
+            reorder=self.reorder,
+            out=out,
+        )
+
+
+class Bz2(Codec):
+    """Bz2 codec for numcodecs."""
+
+    codec_id = 'imagecodecs_bz2'
+
+    def __init__(self, level=None):
+        self.level = level
+
+    def encode(self, buf):
+        return imagecodecs.bz2_encode(buf, level=self.level)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.bz2_decode(buf, out=_flat(out))
+
+
+class Cms(Codec):
+    """CMS codec for numcodecs."""
+
+    codec_id = 'imagecodecs_cms'
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def encode(self, buf, out=None):
+        # return imagecodecs.cms_transform(buf)
+        raise NotImplementedError
+
+    def decode(self, buf, out=None):
+        # return imagecodecs.cms_transform(buf)
+        raise NotImplementedError
+
+
+class Deflate(Codec):
+    """Deflate codec for numcodecs."""
+
+    codec_id = 'imagecodecs_deflate'
+
+    def __init__(self, level=None, raw=False):
+        self.level = level
+        self.raw = bool(raw)
+
+    def encode(self, buf):
+        return imagecodecs.deflate_encode(buf, level=self.level, raw=self.raw)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.deflate_decode(buf, out=_flat(out), raw=self.raw)
+
+
+class Delta(Codec):
+    """Delta codec for numcodecs."""
+
+    codec_id = 'imagecodecs_delta'
+
+    def __init__(self, shape=None, dtype=None, axis=-1, dist=1):
+        self.shape = None if shape is None else tuple(shape)
+        self.dtype = None if dtype is None else numpy.dtype(dtype).str
+        self.axis = axis
+        self.dist = dist
+
+    def encode(self, buf):
+        if self.shape is not None or self.dtype is not None:
+            buf = protective_squeeze(numpy.asarray(buf))
+            assert buf.shape == self.shape
+            assert buf.dtype == self.dtype
+        return imagecodecs.delta_encode(
+            buf, axis=self.axis, dist=self.dist
+        ).tobytes()
+
+    def decode(self, buf, out=None):
+        if self.shape is not None or self.dtype is not None:
+            buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape)
+        return imagecodecs.delta_decode(
+            buf, axis=self.axis, dist=self.dist, out=out
+        )
+
+
+class Float24(Codec):
+    """Float24 codec for numcodecs."""
+
+    codec_id = 'imagecodecs_float24'
+
+    def __init__(self, byteorder=None, rounding=None):
+        self.byteorder = byteorder
+        self.rounding = rounding
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.float24_encode(
+            buf, byteorder=self.byteorder, rounding=self.rounding
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.float24_decode(
+            buf, byteorder=self.byteorder, out=out
+        )
+
+
+class FloatPred(Codec):
+    """Floating Point Predictor codec for numcodecs."""
+
+    codec_id = 'imagecodecs_floatpred'
+
+    def __init__(self, shape, dtype, axis=-1, dist=1):
+        self.shape = tuple(shape)
+        self.dtype = numpy.dtype(dtype).str
+        self.axis = axis
+        self.dist = dist
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        assert buf.shape == self.shape
+        assert buf.dtype == self.dtype
+        return imagecodecs.floatpred_encode(
+            buf, axis=self.axis, dist=self.dist
+        ).tobytes()
+
+    def decode(self, buf, out=None):
+        if not isinstance(buf, numpy.ndarray):
+            buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape)
+        return imagecodecs.floatpred_decode(
+            buf, axis=self.axis, dist=self.dist, out=out
+        )
+
+
+class Gif(Codec):
+    """GIF codec for numcodecs."""
+
+    codec_id = 'imagecodecs_gif'
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.gif_encode(buf)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.gif_decode(buf, asrgb=False, out=out)
+
+
+class Heif(Codec):
+    """HEIF codec for numcodecs."""
+
+    codec_id = 'imagecodecs_heif'
+
+    def __init__(
+        self,
+        level=None,
+        bitspersample=None,
+        photometric=None,
+        compression=None,
+        numthreads=None,
+        index=None,
+    ):
+        self.level = level
+        self.bitspersample = bitspersample
+        self.photometric = photometric
+        self.compression = compression
+        self.numthreads = numthreads
+        self.index = index
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.heif_encode(
+            buf,
+            level=self.level,
+            bitspersample=self.bitspersample,
+            photometric=self.photometric,
+            compression=self.compression,
+            numthreads=self.numthreads,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.heif_decode(
+            buf,
+            index=self.index,
+            photometric=self.photometric,
+            numthreads=self.numthreads,
+            out=out,
+        )
+
+
+class Jetraw(Codec):
+    """Jetraw codec for numcodecs."""
+
+    codec_id = 'imagecodecs_jetraw'
+
+    def __init__(
+        self,
+        shape,
+        identifier,
+        parameters=None,
+        verbosity=None,
+        errorbound=None,
+    ):
+        self.shape = shape
+        self.identifier = identifier
+        self.errorbound = errorbound
+        imagecodecs.jetraw_init(parameters, verbosity)
+
+    def encode(self, buf):
+        return imagecodecs.jetraw_encode(
+            buf, identifier=self.identifier, errorbound=self.errorbound
+        )
+
+    def decode(self, buf, out=None):
+        if out is None:
+            out = numpy.empty(self.shape, numpy.uint16)
+        return imagecodecs.jetraw_decode(buf, out=out)
+
+
+class Jpeg(Codec):
+    """JPEG codec for numcodecs."""
+
+    codec_id = 'imagecodecs_jpeg'
+
+    def __init__(
+        self,
+        bitspersample=None,
+        tables=None,
+        header=None,
+        colorspace_data=None,
+        colorspace_jpeg=None,
+        level=None,
+        subsampling=None,
+        optimize=None,
+        smoothing=None,
+    ):
+        self.tables = tables
+        self.header = header
+        self.bitspersample = bitspersample
+        self.colorspace_data = colorspace_data
+        self.colorspace_jpeg = colorspace_jpeg
+        self.level = level
+        self.subsampling = subsampling
+        self.optimize = optimize
+        self.smoothing = smoothing
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.jpeg_encode(
+            buf,
+            level=self.level,
+            colorspace=self.colorspace_data,
+            outcolorspace=self.colorspace_jpeg,
+            subsampling=self.subsampling,
+            optimize=self.optimize,
+            smoothing=self.smoothing,
+        )
+
+    def decode(self, buf, out=None):
+        out_shape = None
+        if out is not None:
+            out_shape = out.shape
+            out = protective_squeeze(out)
+        img = imagecodecs.jpeg_decode(
+            buf,
+            bitspersample=self.bitspersample,
+            tables=self.tables,
+            header=self.header,
+            colorspace=self.colorspace_jpeg,
+            outcolorspace=self.colorspace_data,
+            out=out,
+        )
+        if out_shape is not None:
+            img = img.reshape(out_shape)
+        return img
+
+    def get_config(self):
+        """Return dictionary holding configuration parameters."""
+        config = dict(id=self.codec_id)
+        for key in self.__dict__:
+            if not key.startswith('_'):
+                value = getattr(self, key)
+                if value is not None and key in ('header', 'tables'):
+                    import base64
+
+                    value = base64.b64encode(value).decode()
+                config[key] = value
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        """Instantiate codec from configuration object."""
+        for key in ('header', 'tables'):
+            value = config.get(key, None)
+            if value is not None and isinstance(value, str):
+                import base64
+
+                config[key] = base64.b64decode(value.encode())
+        return cls(**config)
+
+
+class Jpeg2k(Codec):
+    """JPEG 2000 codec for numcodecs."""
+
+    codec_id = 'imagecodecs_jpeg2k'
+
+    def __init__(
+        self,
+        level=None,
+        codecformat=None,
+        colorspace=None,
+        tile=None,
+        reversible=None,
+        bitspersample=None,
+        resolutions=None,
+        numthreads=None,
+        verbose=0,
+    ):
+        self.level = level
+        self.codecformat = codecformat
+        self.colorspace = colorspace
+        self.tile = None if tile is None else tuple(tile)
+        self.reversible = reversible
+        self.bitspersample = bitspersample
+        self.resolutions = resolutions
+        self.numthreads = numthreads
+        self.verbose = verbose
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.jpeg2k_encode(
+            buf,
+            level=self.level,
+            codecformat=self.codecformat,
+            colorspace=self.colorspace,
+            tile=self.tile,
+            reversible=self.reversible,
+            bitspersample=self.bitspersample,
+            resolutions=self.resolutions,
+            numthreads=self.numthreads,
+            verbose=self.verbose,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.jpeg2k_decode(
+            buf, verbose=self.verbose, numthreads=self.numthreads, out=out
+        )
+
+
+class JpegLs(Codec):
+    """JPEG LS codec for numcodecs."""
+
+    codec_id = 'imagecodecs_jpegls'
+
+    def __init__(self, level=None):
+        self.level = level
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.jpegls_encode(buf, level=self.level)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.jpegls_decode(buf, out=out)
+
+
+class JpegXl(Codec):
+    """JPEG XL codec for numcodecs."""
+
+    codec_id = 'imagecodecs_jpegxl'
+
+    def __init__(
+        self,
+        # encode
+        level=None,
+        effort=None,
+        distance=None,
+        lossless=None,
+        decodingspeed=None,
+        photometric=None,
+        planar=None,
+        usecontainer=None,
+        # decode
+        index=None,
+        keeporientation=None,
+        # both
+        numthreads=None,
+    ):
+        """
+        Return JPEG XL image from numpy array.
+        Float must be in nominal range 0..1.
+
+        Currently L, LA, RGB, RGBA images are supported in contig mode.
+        Extra channels are only supported for grayscale images in planar mode.
+        
+        Parameters
+        ----------
+        level : Default to None, i.e. not overwriting lossess and decodingspeed options.
+            When < 0: Use lossless compression
+            When in [0,1,2,3,4]: Sets the decoding speed tier for the provided options. 
+                Minimum is 0 (slowest to decode, best quality/density), and maximum 
+                is 4 (fastest to decode, at the cost of some quality/density).
+        effort : Default to 3.
+            Sets encoder effort/speed level without affecting decoding speed. 
+            Valid values are, from faster to slower speed: 1:lightning 2:thunder 
+                3:falcon 4:cheetah 5:hare 6:wombat 7:squirrel 8:kitten 9:tortoise. 
+            Speed: lightning, thunder, falcon, cheetah, hare, wombat, squirrel, kitten, tortoise 
+            control the encoder effort in ascending order. 
+            This also affects memory usage: using lower effort will typically reduce memory 
+            consumption during encoding.
+            lightning and thunder are fast modes useful for lossless mode (modular).
+            falcon disables all of the following tools.
+            cheetah enables coefficient reordering, context clustering, and heuristics for selecting DCT sizes and quantization steps.
+            hare enables Gaborish filtering, chroma from luma, and an initial estimate of quantization steps.
+            wombat enables error diffusion quantization and full DCT size selection heuristics.
+            squirrel (default) enables dots, patches, and spline detection, and full context clustering.
+            kitten optimizes the adaptive quantization for a psychovisual metric.
+            tortoise enables a more thorough adaptive quantization search.
+        distance : Default to 1.0
+            Sets the distance level for lossy compression: target max butteraugli distance, 
+            lower = higher quality. Range: 0 .. 15. 0.0 = mathematically lossless 
+            (however, use JxlEncoderSetFrameLossless instead to use true lossless, 
+            as setting distance to 0 alone is not the only requirement). 
+            1.0 = visually lossless. Recommended range: 0.5 .. 3.0.
+        lossess : Default to False. 
+            Use lossess encoding.
+        decodingspeed : Default to 0.
+            Duplicate to level. [0,4]
+        photometric : Return JxlColorSpace value. 
+            Default logic is quite complicated but works most of the time.
+            Accepted value:
+                int: [-1,3]
+                str: ['RGB', 
+                    'WHITEISZERO', 'MINISWHITE', 
+                    'BLACKISZERO', 'MINISBLACK', 'GRAY',
+                    'XYB', 'KNOWN']
+        planar : Enable multi-channel mode.
+            Default to false.
+        usecontainer : 
+            Forces the encoder to use the box-based container format (BMFF) 
+            even when not necessary.
+            When using JxlEncoderUseBoxes, JxlEncoderStoreJPEGMetadata or 
+            JxlEncoderSetCodestreamLevel with level 10, the encoder will 
+            automatically also use the container format, it is not necessary 
+            to use JxlEncoderUseContainer for those use cases.
+            By default this setting is disabled.
+        index : Selectively decode frames for animation.
+            Default to 0, decode all frames.
+            When set to > 0, decode that frame index only.
+        keeporientation : 
+            Enables or disables preserving of as-in-bitstream pixeldata orientation. 
+            Some images are encoded with an Orientation tag indicating that the 
+            decoder must perform a rotation and/or mirroring to the encoded image data.
+
+            If skip_reorientation is JXL_FALSE (the default): the decoder will apply 
+            the transformation from the orientation setting, hence rendering the image 
+            according to its specified intent. When producing a JxlBasicInfo, the decoder 
+            will always set the orientation field to JXL_ORIENT_IDENTITY (matching the 
+            returned pixel data) and also align xsize and ysize so that they correspond 
+            to the width and the height of the returned pixel data.
+
+            If skip_reorientation is JXL_TRUE: the decoder will skip applying the 
+            transformation from the orientation setting, returning the image in 
+            the as-in-bitstream pixeldata orientation. This may be faster to decode 
+            since the decoder doesnt have to apply the transformation, but can 
+            cause wrong display of the image if the orientation tag is not correctly 
+            taken into account by the user.
+
+            By default, this option is disabled, and the returned pixel data is 
+            re-oriented according to the images Orientation setting.
+        threads : Default to 1.
+            If <= 0, use all cores.
+            If > 32, clipped to 32.
+        """
+
+        self.level = level
+        self.effort = effort
+        self.distance = distance
+        self.lossless = bool(lossless)
+        self.decodingspeed = decodingspeed
+        self.photometric = photometric
+        self.planar = planar
+        self.usecontainer = usecontainer
+        self.index = index
+        self.keeporientation = keeporientation
+        self.numthreads = numthreads
+
+    def encode(self, buf):
+        # TODO: only squeeze all but last dim
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.jpegxl_encode(
+            buf,
+            level=self.level,
+            effort=self.effort,
+            distance=self.distance,
+            lossless=self.lossless,
+            decodingspeed=self.decodingspeed,
+            photometric=self.photometric,
+            planar=self.planar,
+            usecontainer=self.usecontainer,
+            numthreads=self.numthreads,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.jpegxl_decode(
+            buf,
+            index=self.index,
+            keeporientation=self.keeporientation,
+            numthreads=self.numthreads,
+            out=out,
+        )
+
+
+class JpegXr(Codec):
+    """JPEG XR codec for numcodecs."""
+
+    codec_id = 'imagecodecs_jpegxr'
+
+    def __init__(
+        self,
+        level=None,
+        photometric=None,
+        hasalpha=None,
+        resolution=None,
+        fp2int=None,
+    ):
+        self.level = level
+        self.photometric = photometric
+        self.hasalpha = hasalpha
+        self.resolution = resolution
+        self.fp2int = fp2int
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.jpegxr_encode(
+            buf,
+            level=self.level,
+            photometric=self.photometric,
+            hasalpha=self.hasalpha,
+            resolution=self.resolution,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.jpegxr_decode(buf, fp2int=self.fp2int, out=out)
+
+
+class Lerc(Codec):
+    """LERC codec for numcodecs."""
+
+    codec_id = 'imagecodecs_lerc'
+
+    def __init__(self, level=None, version=None, planar=None):
+        self.level = level
+        self.version = version
+        self.planar = bool(planar)
+        # TODO: support mask?
+        # self.mask = None
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.lerc_encode(
+            buf,
+            level=self.level,
+            version=self.version,
+            planar=self.planar,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.lerc_decode(buf, out=out)
+
+
+class Ljpeg(Codec):
+    """LJPEG codec for numcodecs."""
+
+    codec_id = 'imagecodecs_ljpeg'
+
+    def __init__(self, bitspersample=None):
+        self.bitspersample = bitspersample
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.ljpeg_encode(buf, bitspersample=self.bitspersample)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.ljpeg_decode(buf, out=out)
+
+
+class Lz4(Codec):
+    """LZ4 codec for numcodecs."""
+
+    codec_id = 'imagecodecs_lz4'
+
+    def __init__(self, level=None, hc=False, header=True):
+        self.level = level
+        self.hc = hc
+        self.header = bool(header)
+
+    def encode(self, buf):
+        return imagecodecs.lz4_encode(
+            buf, level=self.level, hc=self.hc, header=self.header
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.lz4_decode(buf, header=self.header, out=_flat(out))
+
+
+class Lz4f(Codec):
+    """LZ4F codec for numcodecs."""
+
+    codec_id = 'imagecodecs_lz4f'
+
+    def __init__(
+        self,
+        level=None,
+        blocksizeid=False,
+        contentchecksum=None,
+        blockchecksum=None,
+    ):
+        self.level = level
+        self.blocksizeid = blocksizeid
+        self.contentchecksum = contentchecksum
+        self.blockchecksum = blockchecksum
+
+    def encode(self, buf):
+        return imagecodecs.lz4f_encode(
+            buf,
+            level=self.level,
+            blocksizeid=self.blocksizeid,
+            contentchecksum=self.contentchecksum,
+            blockchecksum=self.blockchecksum,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.lz4f_decode(buf, out=_flat(out))
+
+
+class Lzf(Codec):
+    """LZF codec for numcodecs."""
+
+    codec_id = 'imagecodecs_lzf'
+
+    def __init__(self, header=True):
+        self.header = bool(header)
+
+    def encode(self, buf):
+        return imagecodecs.lzf_encode(buf, header=self.header)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.lzf_decode(buf, header=self.header, out=_flat(out))
+
+
+class Lzma(Codec):
+    """LZMA codec for numcodecs."""
+
+    codec_id = 'imagecodecs_lzma'
+
+    def __init__(self, level=None):
+        self.level = level
+
+    def encode(self, buf):
+        return imagecodecs.lzma_encode(buf, level=self.level)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.lzma_decode(buf, out=_flat(out))
+
+
+class Lzw(Codec):
+    """LZW codec for numcodecs."""
+
+    codec_id = 'imagecodecs_lzw'
+
+    def encode(self, buf):
+        return imagecodecs.lzw_encode(buf)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.lzw_decode(buf, out=_flat(out))
+
+
+class PackBits(Codec):
+    """PackBits codec for numcodecs."""
+
+    codec_id = 'imagecodecs_packbits'
+
+    def __init__(self, axis=None):
+        self.axis = axis
+
+    def encode(self, buf):
+        if not isinstance(buf, (bytes, bytearray)):
+            buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.packbits_encode(buf, axis=self.axis)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.packbits_decode(buf, out=_flat(out))
+
+
+class Pglz(Codec):
+    """PGLZ codec for numcodecs."""
+
+    codec_id = 'imagecodecs_pglz'
+
+    def __init__(self, header=True, strategy=None):
+        self.header = bool(header)
+        self.strategy = strategy
+
+    def encode(self, buf):
+        return imagecodecs.pglz_encode(
+            buf, strategy=self.strategy, header=self.header
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.pglz_decode(buf, header=self.header, out=_flat(out))
+
+
+class Png(Codec):
+    """PNG codec for numcodecs."""
+
+    codec_id = 'imagecodecs_png'
+
+    def __init__(self, level=None):
+        self.level = level
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.png_encode(buf, level=self.level)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.png_decode(buf, out=out)
+
+
+class Qoi(Codec):
+    """QOI codec for numcodecs."""
+
+    codec_id = 'imagecodecs_qoi'
+
+    def __init__(self):
+        pass
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.qoi_encode(buf)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.qoi_decode(buf, out=out)
+
+
+class Rgbe(Codec):
+    """RGBE codec for numcodecs."""
+
+    codec_id = 'imagecodecs_rgbe'
+
+    def __init__(self, header=False, shape=None, rle=None):
+        if not header and shape is None:
+            raise ValueError('must specify data shape if no header')
+        if shape and shape[-1] != 3:
+            raise ValueError('invalid shape')
+        self.shape = shape
+        self.header = bool(header)
+        self.rle = None if rle is None else bool(rle)
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.rgbe_encode(buf, header=self.header, rle=self.rle)
+
+    def decode(self, buf, out=None):
+        if out is None and not self.header:
+            out = numpy.empty(self.shape, numpy.float32)
+        return imagecodecs.rgbe_decode(
+            buf, header=self.header, rle=self.rle, out=out
+        )
+
+
+class Rcomp(Codec):
+    """Rcomp codec for numcodecs."""
+
+    codec_id = 'imagecodecs_rcomp'
+
+    def __init__(self, shape, dtype, nblock=None):
+        self.shape = tuple(shape)
+        self.dtype = numpy.dtype(dtype).str
+        self.nblock = nblock
+
+    def encode(self, buf):
+        return imagecodecs.rcomp_encode(buf, nblock=self.nblock)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.rcomp_decode(
+            buf,
+            shape=self.shape,
+            dtype=self.dtype,
+            nblock=self.nblock,
+            out=out,
+        )
+
+
+class Snappy(Codec):
+    """Snappy codec for numcodecs."""
+
+    codec_id = 'imagecodecs_snappy'
+
+    def encode(self, buf):
+        return imagecodecs.snappy_encode(buf)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.snappy_decode(buf, out=_flat(out))
+
+
+class Spng(Codec):
+    """SPNG codec for numcodecs."""
+
+    codec_id = 'imagecodecs_spng'
+
+    def __init__(self, level=None):
+        self.level = level
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.spng_encode(buf, level=self.level)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.spng_decode(buf, out=out)
+
+
+class Tiff(Codec):
+    """TIFF codec for numcodecs."""
+
+    codec_id = 'imagecodecs_tiff'
+
+    def __init__(self, index=None, asrgb=None, verbose=0):
+        self.index = index
+        self.asrgb = bool(asrgb)
+        self.verbose = verbose
+
+    def encode(self, buf):
+        # TODO: not implemented
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.tiff_encode(buf)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.tiff_decode(
+            buf,
+            index=self.index,
+            asrgb=self.asrgb,
+            verbose=self.verbose,
+            out=out,
+        )
+
+
+class Webp(Codec):
+    """WebP codec for numcodecs."""
+
+    codec_id = 'imagecodecs_webp'
+
+    def __init__(self, level=None, lossless=None, method=None, hasalpha=None):
+        self.level = level
+        self.hasalpha = bool(hasalpha)
+        self.method = method
+        self.lossless = lossless
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.webp_encode(
+            buf, level=self.level, lossless=self.lossless, method=self.method
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.webp_decode(buf, hasalpha=self.hasalpha, out=out)
+
+
+class Xor(Codec):
+    """XOR codec for numcodecs."""
+
+    codec_id = 'imagecodecs_xor'
+
+    def __init__(self, shape=None, dtype=None, axis=-1):
+        self.shape = None if shape is None else tuple(shape)
+        self.dtype = None if dtype is None else numpy.dtype(dtype).str
+        self.axis = axis
+
+    def encode(self, buf):
+        if self.shape is not None or self.dtype is not None:
+            buf = protective_squeeze(numpy.asarray(buf))
+            assert buf.shape == self.shape
+            assert buf.dtype == self.dtype
+        return imagecodecs.xor_encode(buf, axis=self.axis).tobytes()
+
+    def decode(self, buf, out=None):
+        if self.shape is not None or self.dtype is not None:
+            buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape)
+        return imagecodecs.xor_decode(buf, axis=self.axis, out=_flat(out))
+
+
+class Zfp(Codec):
+    """ZFP codec for numcodecs."""
+
+    codec_id = 'imagecodecs_zfp'
+
+    def __init__(
+        self,
+        shape=None,
+        dtype=None,
+        strides=None,
+        level=None,
+        mode=None,
+        execution=None,
+        numthreads=None,
+        chunksize=None,
+        header=True,
+    ):
+        if header:
+            self.shape = None
+            self.dtype = None
+            self.strides = None
+        elif shape is None or dtype is None:
+            raise ValueError('invalid shape or dtype')
+        else:
+            self.shape = tuple(shape)
+            self.dtype = numpy.dtype(dtype).str
+            self.strides = None if strides is None else tuple(strides)
+        self.level = level
+        self.mode = mode
+        self.execution = execution
+        self.numthreads = numthreads
+        self.chunksize = chunksize
+        self.header = bool(header)
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        if not self.header:
+            assert buf.shape == self.shape
+            assert buf.dtype == self.dtype
+        return imagecodecs.zfp_encode(
+            buf,
+            level=self.level,
+            mode=self.mode,
+            execution=self.execution,
+            header=self.header,
+            numthreads=self.numthreads,
+            chunksize=self.chunksize,
+        )
+
+    def decode(self, buf, out=None):
+        if self.header:
+            return imagecodecs.zfp_decode(buf, out=out)
+        return imagecodecs.zfp_decode(
+            buf,
+            shape=self.shape,
+            dtype=numpy.dtype(self.dtype),
+            strides=self.strides,
+            numthreads=self.numthreads,
+            out=out,
+        )
+
+
+class Zlib(Codec):
+    """Zlib codec for numcodecs."""
+
+    codec_id = 'imagecodecs_zlib'
+
+    def __init__(self, level=None):
+        self.level = level
+
+    def encode(self, buf):
+        return imagecodecs.zlib_encode(buf, level=self.level)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.zlib_decode(buf, out=_flat(out))
+
+
+class Zlibng(Codec):
+    """Zlibng codec for numcodecs."""
+
+    codec_id = 'imagecodecs_zlibng'
+
+    def __init__(self, level=None):
+        self.level = level
+
+    def encode(self, buf):
+        return imagecodecs.zlibng_encode(buf, level=self.level)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.zlibng_decode(buf, out=_flat(out))
+
+
+class Zopfli(Codec):
+    """Zopfli codec for numcodecs."""
+
+    codec_id = 'imagecodecs_zopfli'
+
+    def encode(self, buf):
+        return imagecodecs.zopfli_encode(buf)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.zopfli_decode(buf, out=_flat(out))
+
+
+class Zstd(Codec):
+    """ZStandard codec for numcodecs."""
+
+    codec_id = 'imagecodecs_zstd'
+
+    def __init__(self, level=None):
+        self.level = level
+
+    def encode(self, buf):
+        return imagecodecs.zstd_encode(buf, level=self.level)
+
+    def decode(self, buf, out=None):
+        return imagecodecs.zstd_decode(buf, out=_flat(out))
+
+
+def _flat(out):
+    """Return numpy array as contiguous view of bytes if possible."""
+    if out is None:
+        return None
+    view = memoryview(out)
+    if view.readonly or not view.contiguous:
+        return None
+    return view.cast('B')
+
+
+def register_codecs(codecs=None, force=False, verbose=True):
+    """Register codecs in this module with numcodecs."""
+    for name, cls in globals().items():
+        if not hasattr(cls, 'codec_id') or name == 'Codec':
+            continue
+        if codecs is not None and cls.codec_id not in codecs:
+            continue
+        try:
+            try:
+                get_codec({'id': cls.codec_id})
+            except TypeError:
+                # registered, but failed
+                pass
+        except ValueError:
+            # not registered yet
+            pass
+        else:
+            if not force:
+                if verbose:
+                    log_warning(
+                        f'numcodec {cls.codec_id!r} already registered'
+                    )
+                continue
+            if verbose:
+                log_warning(f'replacing registered numcodec {cls.codec_id!r}')
+        register_codec(cls)
+
+
+def log_warning(msg, *args, **kwargs):
+    """Log message with level WARNING."""
+    import logging
+
+    logging.getLogger(__name__).warning(msg, *args, **kwargs)

common/checkpoint_util.py

@@ -0,0 +1,59 @@
+from typing import Optional, Dict
+import os
+
+class TopKCheckpointManager:
+    def __init__(self,
+            save_dir,
+            monitor_key: str,
+            mode='min',
+            k=1,
+            format_str='epoch={epoch:03d}-train_loss={train_loss:.3f}.ckpt'
+        ):
+        assert mode in ['max', 'min']
+        assert k >= 0
+
+        self.save_dir = save_dir
+        self.monitor_key = monitor_key
+        self.mode = mode
+        self.k = k
+        self.format_str = format_str
+        self.path_value_map = dict()
+    
+    def get_ckpt_path(self, data: Dict[str, float]) -> Optional[str]:
+        if self.k == 0:
+            return None
+
+        value = data[self.monitor_key]
+        ckpt_path = os.path.join(
+            self.save_dir, self.format_str.format(**data))
+        
+        if len(self.path_value_map) < self.k:
+            # under-capacity
+            self.path_value_map[ckpt_path] = value
+            return ckpt_path
+        
+        # at capacity
+        sorted_map = sorted(self.path_value_map.items(), key=lambda x: x[1])
+        min_path, min_value = sorted_map[0]
+        max_path, max_value = sorted_map[-1]
+
+        delete_path = None
+        if self.mode == 'max':
+            if value > min_value:
+                delete_path = min_path
+        else:
+            if value < max_value:
+                delete_path = max_path
+
+        if delete_path is None:
+            return None
+        else:
+            del self.path_value_map[delete_path]
+            self.path_value_map[ckpt_path] = value
+
+            if not os.path.exists(self.save_dir):
+                os.mkdir(self.save_dir)
+
+            if os.path.exists(delete_path):
+                os.remove(delete_path)
+            return ckpt_path

common/compute_all_loss.py

@@ -0,0 +1,483 @@
+import torch
+import numpy as np
+import torch.nn.functional as F
+import random
+from einops import rearrange, reduce
+import cv2
+import concurrent.futures
+from diffusion_policy.policy.base_image_policy import BaseImagePolicy
+from diffusion_policy.common.pytorch_util import dict_apply, replace_submodules
+from diffusion_policy.model.common.slice import slice_episode
+
+
+def unflatten_dataset_dict(flat_dict, delimiter='/'):
+    result = {}
+    for compound_key, value in flat_dict.items():
+        keys = compound_key.split(delimiter)
+        current = result
+        for key in keys[:-1]:
+            if key not in current:
+                current[key] = {}
+            current = current[key]
+        current[keys[-1]] = value
+    
+    return result
+
+
+def decode_image(data):
+    return cv2.imdecode(data, 1)
+
+
+def compute_all_traj_loss(replay_buffer=None, model:BaseImagePolicy=None, ref_model:BaseImagePolicy=None, stride=1):
+    if replay_buffer is None:
+        return np.zeros([1])
+    else:
+        data = replay_buffer.data
+        meta_data = replay_buffer.meta
+        observations_1 = np.array(data['obs'], dtype=np.float32)
+        actions_1 = np.array(data['action'], dtype=np.float32)
+        observations_2 = np.array(data['obs_2'], dtype=np.float32)
+        actions_2 = np.array(data['action_2'], dtype=np.float32)
+
+        total_size = len(observations_1)
+
+        # Calculate 25% of the data size
+        sample_size = int(total_size * 0.25)
+
+        # Generate random indices for sampling
+        indices = np.random.choice(total_size, size=sample_size, replace=False)
+
+        # Extract 25% of the data using the indices
+        observations_1 = observations_1[indices]
+        actions_1 = actions_1[indices]
+        observations_2 = observations_2[indices]
+        actions_2 = actions_2[indices]
+
+        for param in ref_model.parameters():
+            param.requires_grad = False
+
+        ref_model = ref_model.to(model.device)
+
+        # Normalize data
+        batch_1 = {
+            'obs': observations_1,
+            'action': actions_1,
+        }
+
+        batch_2 = {
+            'obs': observations_2,
+            'action': actions_2,
+        }
+        nbatch_1 = model.normalizer.normalize(batch_1)
+        nbatch_2 = model.normalizer.normalize(batch_2)
+        obs_1, obs_2 = nbatch_1['obs'], nbatch_2['obs']
+        actions_1, actions_2 = nbatch_1['action'], nbatch_2['action']
+
+        # Slice trajectories
+        obs_1 = slice_episode(obs_1, horizon=model.horizon, stride=stride)
+        action_1 = slice_episode(actions_1, horizon=model.horizon, stride=stride)
+        obs_2 = slice_episode(obs_2, horizon=model.horizon, stride=stride)
+        action_2 = slice_episode(actions_2, horizon=model.horizon, stride=stride)
+
+
+        bsz = obs_1[0].shape[0]
+        timesteps = torch.randint(0, model.noise_scheduler.config.num_train_timesteps, (bsz,), device=model.device).long()
+
+        # Pre-allocate loss
+        traj_loss_1, traj_loss_2 = 0, 0
+
+        # Helper function to compute loss for a single trajectory
+        def compute_traj_loss(obs_slices, action_slices, timestep, model, ref_policy):
+            total_loss = 0
+            
+            for idx, (obs_slide, action_slide) in enumerate(zip(obs_slices, action_slices)):
+                gamma_factors = model.gamma ** (idx * model.horizon + torch.arange(model.horizon, device=model.device))
+                if model.obs_as_cond:
+                    cond = obs_slide[:, :model.n_obs_steps, :]
+                    cond = cond.detach().to(model.device)
+                    # cond.detach().to(model.device)
+                    trajectory = action_slide[:, -model.n_action_steps:] if model.pred_action_steps_only else action_slide
+                else:
+                    cond = None
+                    trajectory = np.concatenate([action_slide, obs_slide], axis=-1)
+
+                condition_mask = model.mask_generator(trajectory.shape).to(model.device)
+                loss_mask = (~condition_mask).float()
+
+                trajectory = torch.tensor(trajectory, device=model.device, dtype=torch.float32)
+                noise = torch.randn(trajectory.shape, device=model.device)
+
+                # Disable gradient computation
+                with torch.no_grad():
+                    noisy_trajectory = model.noise_scheduler.add_noise(trajectory, noise, timestep)
+                    noisy_trajectory[condition_mask] = trajectory[condition_mask]
+
+                    pred_ref = ref_policy(noisy_trajectory, timestep, cond)
+                    pred = model.model(noisy_trajectory, timestep, cond)
+
+                    pred_type = model.noise_scheduler.config.prediction_type 
+                    if pred_type == 'epsilon':
+                        target = noise
+                    elif pred_type == 'sample':
+                        target = trajectory
+                    else:
+                        raise ValueError(f"Unsupported prediction type {pred_type}")
+                    
+                    loss = F.mse_loss(pred, target, reduction='none')
+                    loss_ref = F.mse_loss(pred_ref, target, reduction='none')
+                    loss = loss * loss_mask.type(loss.dtype)
+                    loss_ref = loss_ref * loss_mask.type(loss.dtype)
+                    loss = reduce(loss, 'b t ... -> b t (...)', 'mean')
+                    loss_ref = reduce(loss_ref, 'b t ... -> b t (...)', 'mean')
+                    
+                    slice_loss = torch.sum((loss - loss_ref), dim=-1)
+                    total_loss += torch.sum(slice_loss * gamma_factors)
+                # Explicitly delete unused variables to release GPU memory
+                del trajectory, noise, noisy_trajectory, pred_ref, pred, loss_mask, condition_mask
+                torch.cuda.empty_cache()
+
+            return total_loss.detach()
+
+        # Compute loss for trajectory 1
+        traj_loss_1 = compute_traj_loss(obs_1, action_1, timesteps, model, ref_model)
+        # Compute loss for trajectory 2
+        traj_loss_2 = compute_traj_loss(obs_2, action_2, timesteps, model, ref_model)
+
+        # Average the losses
+        loss = (traj_loss_1 + traj_loss_2) / 2
+
+        return torch.mean(loss)
+    
+
+def compute_all_traj_loss_realrobot(replay_buffer=None, model=None, ref_model=None, stride=1, sample_size = 20, batch_size=10, gc_every_n_batches=12):
+    import gc
+    gc.collect()
+    torch.cuda.empty_cache()
+    
+    if replay_buffer is None:
+        return np.zeros([1])
+    else:
+        assert sample_size >= batch_size, "data_size should be greater than or equal to batch_size"
+        if hasattr(torch.cuda, 'set_per_process_memory_fraction'):
+            torch.cuda.set_per_process_memory_fraction(0.8)
+        
+        data = replay_buffer.data
+        data = unflatten_dataset_dict(flat_dict=data)
+        data_size = len(data['action'])
+        indicis = np.random.choice(data_size, size=sample_size, replace=False)
+
+        observations_1 = data['obs']
+        actions_1 = np.array(data['action'][indicis], dtype=np.float32)
+        observations_2 = data['obs_2']
+        actions_2 = np.array(data['action_2'][indicis], dtype=np.float32)
+        compress_len_1 = data['compress_len'][indicis]
+        compress_len_2 = data['compress_len_2'][indicis]
+        camera_keys = observations_1['images'].keys()
+        qpos_keys = [key for key in observations_1.keys() if key != 'images']
+        del data
+
+        for key in camera_keys:
+            img_data_1 = observations_1['images'][key][indicis]
+            img_data_2 = observations_2['images'][key][indicis]
+            total_images = img_data_1.shape[0]
+            
+            img_batch_size = min(batch_size, total_images)
+            decompressed_images_1 = []
+            
+            for batch_idx in range(0, total_images, img_batch_size):
+                end_idx = min(batch_idx + img_batch_size, total_images)
+                batch_decompressed = []
+                
+                for k in range(batch_idx, end_idx):
+                    image = img_data_1[k, :, :int(compress_len_1[k, 0])].copy()
+                    with concurrent.futures.ThreadPoolExecutor() as executor:
+                        results = executor.map(decode_image, image)
+                        decompressed_images = list(results)
+                    batch_decompressed.append(decompressed_images)
+                
+                batch_decompressed = np.array(batch_decompressed)
+                batch_decompressed = np.einsum('b k h w c -> b k c h w', batch_decompressed)
+                decompressed_images_1.append(torch.from_numpy(batch_decompressed / 255.0).float())
+                
+                del batch_decompressed
+                torch.cuda.empty_cache()
+                
+                if batch_idx % (img_batch_size * gc_every_n_batches) == 0:
+                    gc.collect()
+            
+            observations_1[key] = torch.cat(decompressed_images_1, dim=0)
+            del observations_1['images']
+            del decompressed_images_1
+            
+            decompressed_images_2 = []
+            
+            for batch_idx in range(0, total_images, img_batch_size):
+                end_idx = min(batch_idx + img_batch_size, total_images)
+                batch_decompressed = []
+                
+                for k in range(batch_idx, end_idx):
+                    image = img_data_2[k, :, :int(compress_len_2[k, 0])].copy()
+                    with concurrent.futures.ThreadPoolExecutor() as executor:
+                        results = executor.map(decode_image, image)
+                        decompressed_images = list(results)
+                    batch_decompressed.append(decompressed_images)
+                
+                batch_decompressed = np.array(batch_decompressed)
+                batch_decompressed = np.einsum('b k h w c -> b k c h w', batch_decompressed)
+                decompressed_images_2.append(torch.from_numpy(batch_decompressed / 255.0).float())
+                
+                del batch_decompressed
+                torch.cuda.empty_cache()
+                
+                if batch_idx % (img_batch_size * gc_every_n_batches) == 0:
+                    gc.collect()
+            
+            observations_2[key] = torch.cat(decompressed_images_2, dim=0)
+            del observations_2['images']
+            del decompressed_images_2
+            
+            del img_data_1, img_data_2
+            gc.collect()
+            torch.cuda.empty_cache()
+
+        for key in qpos_keys:
+            observations_1[key] = torch.from_numpy(observations_1[key]).float()
+            observations_2[key] = torch.from_numpy(observations_2[key]).float()
+            
+            observations_1[key] = observations_1[key].cpu()
+            observations_2[key] = observations_2[key].cpu()
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        for param in ref_model.parameters():
+            param.requires_grad = False
+        
+        device = model.device
+        ref_model = ref_model.to(device)
+        
+        with torch.no_grad():
+            obs_1 = model.normalizer.normalize(observations_1)
+            action_1 = model.normalizer['action'].normalize(actions_1)
+            obs_2 = model.normalizer.normalize(observations_2)
+            action_2 = model.normalizer['action'].normalize(actions_2)
+        
+        start_1 = random.randint(0, model.n_obs_steps)
+        start_2 = random.randint(0, model.n_obs_steps)
+        
+        with torch.no_grad():
+            obs_1 = {key: slice_episode(obs_1[key], horizon=model.horizon, stride=stride, start=start_1) for key in obs_1.keys()}
+            action_1 = slice_episode(action_1, horizon=model.horizon, stride=stride, start=start_1)
+            obs_2 = {key: slice_episode(obs_2[key], horizon=model.horizon, stride=stride, start=start_2) for key in obs_2.keys()}
+            action_2 = slice_episode(action_2, horizon=model.horizon, stride=stride, start=start_2)
+        
+        del observations_1, observations_2, actions_1, actions_2
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        def compute_traj_image_loss_batched(obs_slices, action_slices, model, ref_model, batch_size):
+            with torch.no_grad(): 
+                To = model.n_obs_steps
+                horizon = model.horizon
+                total_samples = action_slices.shape[0]
+                total_loss = torch.zeros(total_samples, device='cpu')
+                
+                num_batches = (total_samples + batch_size - 1) // batch_size
+                device = model.device
+                
+                for batch_idx in range(num_batches):
+
+                    if batch_idx > 0 and batch_idx % gc_every_n_batches == 0:
+                        gc.collect()
+                        torch.cuda.empty_cache()
+                    
+                    start_idx = batch_idx * batch_size
+                    end_idx = min(start_idx + batch_size, total_samples)
+                    current_batch_size = end_idx - start_idx
+                    
+                    batch_timesteps = torch.randint(0, model.noise_scheduler.config.num_train_timesteps, 
+                                                   (current_batch_size,), device=device).long()
+                    
+                    batch_action_slices = action_slices[start_idx:end_idx].to(device)
+                    batch_obs_slices = {key: obs_slices[key][start_idx:end_idx].to(device) for key in obs_slices.keys()}
+                    
+                    batch_total_loss = torch.zeros(current_batch_size, device=device)
+                    
+                    for idx in range(current_batch_size):
+                        action_slide = batch_action_slices[idx:idx+1]
+                        obs_slide = {key: batch_obs_slices[key][idx:idx+1] for key in batch_obs_slices.keys()}
+                        sample_timesteps = batch_timesteps[idx:idx+1]
+                        
+                        local_cond = None
+                        global_cond = None
+                        global_cond_ref = None
+
+                        if model.obs_as_global_cond:
+                            this_nobs = dict_apply(obs_slide, 
+                                lambda x: x[:,:To,...].reshape(-1, *x.shape[2:]))
+                            nobs_features = model.obs_encoder(this_nobs)
+                            nobs_features_ref = ref_model.obs_encoder(this_nobs)
+                            
+                            global_cond = nobs_features.reshape(1, -1)
+                            global_cond_ref = nobs_features_ref.reshape(1, -1)
+                            trajectory = action_slide
+                        else:
+                            this_nobs = dict_apply(obs_slide, 
+                                lambda x: x.reshape(-1, *x.shape[2:]))
+                            nobs_features = model.obs_encoder(this_nobs)
+                            nobs_features_ref = ref_model.obs_encoder(this_nobs)
+                            
+                            nobs_features = nobs_features.reshape(1, horizon, -1)
+                            nobs_features_ref = nobs_features_ref.reshape(1, horizon, -1)
+                            
+                            trajectory = torch.cat([action_slide, nobs_features], dim=-1)
+                            trajectory_ref = torch.cat([action_slide, nobs_features_ref], dim=-1)
+
+                        condition_mask = model.mask_generator(trajectory.shape).to(device)
+                        loss_mask = (~condition_mask).float()
+
+                        noise = torch.randn(trajectory.shape, device=device)
+
+                        noisy_trajectory = model.noise_scheduler.add_noise(trajectory, noise, sample_timesteps)
+                        noisy_trajectory[condition_mask] = trajectory[condition_mask]
+                        
+                        if not model.obs_as_global_cond:
+                            noisy_trajectory_ref = model.noise_scheduler.add_noise(trajectory_ref, noise, sample_timesteps)
+                            noisy_trajectory_ref[condition_mask] = trajectory_ref[condition_mask]
+                        else:
+                            noisy_trajectory_ref = noisy_trajectory.clone()
+
+                        pred = model.model(noisy_trajectory, sample_timesteps, 
+                                          local_cond=local_cond, global_cond=global_cond)
+                        pred_ref = ref_model.model(noisy_trajectory_ref if not model.obs_as_global_cond else noisy_trajectory, 
+                                                 sample_timesteps, local_cond=local_cond, global_cond=global_cond_ref)
+
+                        pred_type = model.noise_scheduler.config.prediction_type
+                        if pred_type == 'epsilon':
+                            target = noise
+                        elif pred_type == 'sample':
+                            target = trajectory
+                        else:
+                            raise ValueError(f"Unsupported prediction type {pred_type}")
+
+                        loss = F.mse_loss(pred, target, reduction='none')
+                        loss_ref = F.mse_loss(pred_ref, target, reduction='none')
+                        
+                        loss = loss * loss_mask
+                        loss_ref = loss_ref * loss_mask
+                        loss = reduce(loss, 'b t ... -> b t (...)', 'mean')
+                        loss_ref = reduce(loss_ref, 'b t ... -> b t (...)', 'mean')
+                        
+                        slice_loss = torch.sum(loss_ref - loss, dim=1)
+                        batch_total_loss[idx] = slice_loss.squeeze()
+                        
+                        del trajectory, noise, noisy_trajectory, pred, pred_ref
+                        if not model.obs_as_global_cond:
+                            del trajectory_ref, noisy_trajectory_ref
+                        del nobs_features, nobs_features_ref, this_nobs
+                    
+                    total_loss[start_idx:end_idx] = batch_total_loss.cpu()
+                    
+                    del batch_action_slices, batch_obs_slices, batch_timesteps, batch_total_loss
+                    torch.cuda.empty_cache()
+
+                return total_loss
+
+        with torch.no_grad():
+            total_samples = action_1.shape[0]
+            traj_loss_1 = compute_traj_image_loss_batched(obs_1, action_1, model, ref_model, batch_size)
+            
+            if traj_loss_1.device.type != 'cpu':
+                traj_loss_1 = traj_loss_1.cpu()
+            
+            del obs_1, action_1
+            gc.collect()
+            torch.cuda.empty_cache()
+            
+            total_samples = action_2.shape[0]
+            traj_loss_2 = compute_traj_image_loss_batched(obs_2, action_2, model, ref_model, batch_size)
+            
+            if traj_loss_2.device.type != 'cpu':
+                traj_loss_2 = traj_loss_2.cpu()
+            
+            loss = (traj_loss_1 + traj_loss_2) / 2
+            final_loss = torch.mean(loss)
+
+        del obs_2, action_2, traj_loss_1, traj_loss_2, loss
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        return final_loss
+
+    
+def compute_all_bet_traj_loss(replay_buffer=None, model=None, stride=1):
+    if replay_buffer is None:
+        return np.zeros([1])
+    else:
+        data = replay_buffer.data
+        meta_data = replay_buffer.meta
+        observations_1 = np.array(data['obs'], dtype=np.float32)
+        actions_1 = np.array(data['action'], dtype=np.float32)
+        observations_2 = np.array(data['obs_2'], dtype=np.float32)
+        actions_2 = np.array(data['action_2'], dtype=np.float32)
+        length_1 = torch.tensor(meta_data['length'], device=model.device)
+        length_2 = torch.tensor(meta_data['length_2'], device=model.device)
+
+        # Normalize data
+        batch_1 = {
+            'obs': observations_1,
+            'action': actions_1,
+        }
+
+        batch_2 = {
+            'obs': observations_2,
+            'action': actions_2,
+        }
+        nbatch_1 = model.normalizer.normalize(batch_1)
+        nbatch_2 = model.normalizer.normalize(batch_2)
+        obs_1, obs_2 = nbatch_1['obs'], nbatch_2['obs']
+        actions_1, actions_2 = nbatch_1['action'], nbatch_2['action']
+
+        # Slice trajectories
+        obs_1 = slice_episode(obs_1, horizon=model.horizon, stride=stride)
+        action_1 = slice_episode(actions_1, horizon=model.horizon, stride=stride)
+        obs_2 = slice_episode(obs_2, horizon=model.horizon, stride=stride)
+        action_2 = slice_episode(actions_2, horizon=model.horizon, stride=stride)
+
+        # Pre-allocate loss
+        traj_loss_1, traj_loss_2 = 0, 0
+
+        # Helper function to compute loss for a single trajectory
+        def compute_traj_loss(obs_slices, action_slices, model, length, stride):
+            total_loss = 0
+            
+            for idx, (obs_slide, action_slide) in enumerate(zip(obs_slices, action_slices)):
+                gamma_factors = model.gamma ** (idx * model.horizon)
+                obs_slide[:, model.n_obs_steps:, :] = -2
+
+                enc_obs = model.obs_encoding_net(obs_slide)
+                latent = model.action_ae.encode_into_latent(action_slide, enc_obs)
+
+                loss = model.get_pred_loss(
+                    obs_rep=enc_obs.clone(),
+                    target_latents=latent,
+                )
+
+                mask = (model.horizon + (idx - 1)*stride) <= length
+                mask = mask.int()
+
+                total_loss += (loss * mask) * gamma_factors
+
+            total_loss = torch.sum(total_loss, dim=-1)
+
+            return total_loss.detach()
+
+        # Compute loss for trajectory 1
+        traj_loss_1 = compute_traj_loss(obs_1, action_1, model, length_1, stride)
+        # Compute loss for trajectory 2
+        traj_loss_2 = compute_traj_loss(obs_2, action_2, model, length_2, stride)
+
+        # Average the losses
+        loss = (traj_loss_1 + traj_loss_2) / 2
+
+        return torch.mean(loss)

common/cv2_util.py

@@ -0,0 +1,150 @@
+from typing import Tuple
+import math
+import cv2
+import numpy as np
+
+def draw_reticle(img, u, v, label_color):
+    """
+    Draws a reticle (cross-hair) on the image at the given position on top of
+    the original image.
+    @param img (In/Out) uint8 3 channel image
+    @param u X coordinate (width)
+    @param v Y coordinate (height)
+    @param label_color tuple of 3 ints for RGB color used for drawing.
+    """
+    # Cast to int.
+    u = int(u)
+    v = int(v)
+
+    white = (255, 255, 255)
+    cv2.circle(img, (u, v), 10, label_color, 1)
+    cv2.circle(img, (u, v), 11, white, 1)
+    cv2.circle(img, (u, v), 12, label_color, 1)
+    cv2.line(img, (u, v + 1), (u, v + 3), white, 1)
+    cv2.line(img, (u + 1, v), (u + 3, v), white, 1)
+    cv2.line(img, (u, v - 1), (u, v - 3), white, 1)
+    cv2.line(img, (u - 1, v), (u - 3, v), white, 1)
+
+
+def draw_text(
+    img,
+    *,
+    text,
+    uv_top_left,
+    color=(255, 255, 255),
+    fontScale=0.5,
+    thickness=1,
+    fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+    outline_color=(0, 0, 0),
+    line_spacing=1.5,
+):
+    """
+    Draws multiline with an outline.
+    """
+    assert isinstance(text, str)
+
+    uv_top_left = np.array(uv_top_left, dtype=float)
+    assert uv_top_left.shape == (2,)
+
+    for line in text.splitlines():
+        (w, h), _ = cv2.getTextSize(
+            text=line,
+            fontFace=fontFace,
+            fontScale=fontScale,
+            thickness=thickness,
+        )
+        uv_bottom_left_i = uv_top_left + [0, h]
+        org = tuple(uv_bottom_left_i.astype(int))
+
+        if outline_color is not None:
+            cv2.putText(
+                img,
+                text=line,
+                org=org,
+                fontFace=fontFace,
+                fontScale=fontScale,
+                color=outline_color,
+                thickness=thickness * 3,
+                lineType=cv2.LINE_AA,
+            )
+        cv2.putText(
+            img,
+            text=line,
+            org=org,
+            fontFace=fontFace,
+            fontScale=fontScale,
+            color=color,
+            thickness=thickness,
+            lineType=cv2.LINE_AA,
+        )
+
+        uv_top_left += [0, h * line_spacing]
+
+
+def get_image_transform(
+        input_res: Tuple[int,int]=(1280,720), 
+        output_res: Tuple[int,int]=(640,480), 
+        bgr_to_rgb: bool=False):
+
+    iw, ih = input_res
+    ow, oh = output_res
+    rw, rh = None, None
+    interp_method = cv2.INTER_AREA
+
+    if (iw/ih) >= (ow/oh):
+        # input is wider
+        rh = oh
+        rw = math.ceil(rh / ih * iw)
+        if oh > ih:
+            interp_method = cv2.INTER_LINEAR
+    else:
+        rw = ow
+        rh = math.ceil(rw / iw * ih)
+        if ow > iw:
+            interp_method = cv2.INTER_LINEAR
+    
+    w_slice_start = (rw - ow) // 2
+    w_slice = slice(w_slice_start, w_slice_start + ow)
+    h_slice_start = (rh - oh) // 2
+    h_slice = slice(h_slice_start, h_slice_start + oh)
+    c_slice = slice(None)
+    if bgr_to_rgb:
+        c_slice = slice(None, None, -1)
+
+    def transform(img: np.ndarray):
+        assert img.shape == ((ih,iw,3))
+        # resize
+        img = cv2.resize(img, (rw, rh), interpolation=interp_method)
+        # crop
+        img = img[h_slice, w_slice, c_slice]
+        return img
+    return transform
+
+def optimal_row_cols(
+        n_cameras,
+        in_wh_ratio,
+        max_resolution=(1920, 1080)
+    ):
+    out_w, out_h = max_resolution
+    out_wh_ratio = out_w / out_h
+    
+    n_rows = np.arange(n_cameras,dtype=np.int64) + 1
+    n_cols = np.ceil(n_cameras / n_rows).astype(np.int64)
+    cat_wh_ratio = in_wh_ratio * (n_cols / n_rows)
+    ratio_diff = np.abs(out_wh_ratio - cat_wh_ratio)
+    best_idx = np.argmin(ratio_diff)
+    best_n_row = n_rows[best_idx]
+    best_n_col = n_cols[best_idx]
+    best_cat_wh_ratio = cat_wh_ratio[best_idx]
+
+    rw, rh = None, None
+    if best_cat_wh_ratio >= out_wh_ratio:
+        # cat is wider
+        rw = math.floor(out_w / best_n_col)
+        rh = math.floor(rw / in_wh_ratio)
+    else:
+        rh = math.floor(out_h / best_n_row)
+        rw = math.floor(rh * in_wh_ratio)
+    
+    # crop_resolution = (rw, rh)
+    return rw, rh, best_n_col, best_n_row

common/env_util.py

@@ -0,0 +1,23 @@
+import cv2
+import numpy as np
+
+
+def render_env_video(env, states, actions=None):
+    observations = states
+    imgs = list()
+    for i in range(len(observations)):
+        state = observations[i]
+        env.set_state(state)
+        if i == 0:
+            env.set_state(state)
+        img = env.render()
+        # draw action
+        if actions is not None:
+            action = actions[i]
+            coord = (action / 512 * 96).astype(np.int32)
+            cv2.drawMarker(img, coord, 
+                color=(255,0,0), markerType=cv2.MARKER_CROSS,
+                markerSize=8, thickness=1)
+        imgs.append(img)
+    imgs = np.array(imgs)
+    return imgs

common/json_logger.py

@@ -0,0 +1,117 @@
+from typing import Optional, Callable, Any, Sequence
+import os
+import copy
+import json
+import numbers
+import pandas as pd
+
+
+def read_json_log(path: str, 
+        required_keys: Sequence[str]=tuple(), 
+        **kwargs) -> pd.DataFrame:
+    """
+    Read json-per-line file, with potentially incomplete lines.
+    kwargs passed to pd.read_json
+    """
+    lines = list()
+    with open(path, 'r') as f:
+        while True:
+            # one json per line
+            line = f.readline()
+            if len(line) == 0:
+                # EOF
+                break
+            elif not line.endswith('\n'):
+                # incomplete line
+                break
+            is_relevant = False
+            for k in required_keys:
+                if k in line:
+                    is_relevant = True
+                    break
+            if is_relevant:
+                lines.append(line)
+    if len(lines) < 1:
+        return pd.DataFrame()  
+    json_buf = f'[{",".join([line for line in (line.strip() for line in lines) if line])}]'
+    df = pd.read_json(json_buf, **kwargs)
+    return df
+
+class JsonLogger:
+    def __init__(self, path: str, 
+            filter_fn: Optional[Callable[[str,Any],bool]]=None):
+        if filter_fn is None:
+            filter_fn = lambda k,v: isinstance(v, numbers.Number)
+
+        # default to append mode
+        self.path = path
+        self.filter_fn = filter_fn
+        self.file = None
+        self.last_log = None
+    
+    def start(self):
+        # use line buffering
+        try:
+            self.file = file = open(self.path, 'r+', buffering=1)
+        except FileNotFoundError:
+            self.file = file = open(self.path, 'w+', buffering=1)
+
+        # Move the pointer (similar to a cursor in a text editor) to the end of the file
+        pos = file.seek(0, os.SEEK_END)
+
+        # Read each character in the file one at a time from the last
+        # character going backwards, searching for a newline character
+        # If we find a new line, exit the search
+        while pos > 0 and file.read(1) != "\n":
+            pos -= 1
+            file.seek(pos, os.SEEK_SET)
+        # now the file pointer is at one past the last '\n'
+        # and pos is at the last '\n'.
+        last_line_end = file.tell()
+        
+        # find the start of second last line
+        pos = max(0, pos-1)
+        file.seek(pos, os.SEEK_SET)
+        while pos > 0 and file.read(1) != "\n":
+            pos -= 1
+            file.seek(pos, os.SEEK_SET)
+        # now the file pointer is at one past the second last '\n'
+        last_line_start = file.tell()
+
+        if last_line_start < last_line_end:
+            # has last line of json
+            last_line = file.readline()
+            self.last_log = json.loads(last_line)
+        
+        # remove the last incomplete line
+        file.seek(last_line_end)
+        file.truncate()
+    
+    def stop(self):
+        self.file.close()
+        self.file = None
+    
+    def __enter__(self):
+        self.start()
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()
+    
+    def log(self, data: dict):
+        filtered_data = dict(
+            filter(lambda x: self.filter_fn(*x), data.items()))
+        # save current as last log
+        self.last_log = filtered_data
+        for k, v in filtered_data.items():
+            if isinstance(v, numbers.Integral):
+                filtered_data[k] = int(v)
+            elif isinstance(v, numbers.Number):
+                filtered_data[k] = float(v)
+        buf = json.dumps(filtered_data)
+        # ensure one line per json
+        buf = buf.replace('\n','') + '\n'
+        self.file.write(buf)
+    
+    def get_last_log(self):
+        return copy.deepcopy(self.last_log)

common/nested_dict_util.py

@@ -0,0 +1,32 @@
+import functools
+
+def nested_dict_map(f, x):
+    """
+    Map f over all leaf of nested dict x
+    """
+
+    if not isinstance(x, dict):
+        return f(x)
+    y = dict()
+    for key, value in x.items():
+        y[key] = nested_dict_map(f, value)
+    return y
+
+def nested_dict_reduce(f, x):
+    """
+    Map f over all values of nested dict x, and reduce to a single value
+    """
+    if not isinstance(x, dict):
+        return x
+
+    reduced_values = list()
+    for value in x.values():
+        reduced_values.append(nested_dict_reduce(f, value))
+    y = functools.reduce(f, reduced_values)
+    return y
+
+
+def nested_dict_check(f, x):
+    bool_dict = nested_dict_map(f, x)
+    result = nested_dict_reduce(lambda x, y: x and y, bool_dict)
+    return result

common/normalize_util.py

@@ -0,0 +1,223 @@
+from diffusion_policy.model.common.normalizer import SingleFieldLinearNormalizer
+from diffusion_policy.common.pytorch_util import dict_apply, dict_apply_reduce, dict_apply_split
+import numpy as np
+
+
+def get_range_normalizer_from_stat(stat, output_max=1, output_min=-1, range_eps=1e-7):
+    # -1, 1 normalization
+    input_max = stat['max']
+    input_min = stat['min']
+    input_range = input_max - input_min
+    ignore_dim = input_range < range_eps
+    input_range[ignore_dim] = output_max - output_min
+    scale = (output_max - output_min) / input_range
+    offset = output_min - scale * input_min
+    offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim]
+
+    return SingleFieldLinearNormalizer.create_manual(
+        scale=scale,
+        offset=offset,
+        input_stats_dict=stat
+    )
+
+def get_image_range_normalizer():
+    scale = np.array([2], dtype=np.float32)
+    offset = np.array([-1], dtype=np.float32)
+    stat = {
+        'min': np.array([0], dtype=np.float32),
+        'max': np.array([1], dtype=np.float32),
+        'mean': np.array([0.5], dtype=np.float32),
+        'std': np.array([np.sqrt(1/12)], dtype=np.float32)
+    }
+    return SingleFieldLinearNormalizer.create_manual(
+        scale=scale,
+        offset=offset,
+        input_stats_dict=stat
+    )
+
+def get_identity_normalizer_from_stat(stat):
+    scale = np.ones_like(stat['min'])
+    offset = np.zeros_like(stat['min'])
+    return SingleFieldLinearNormalizer.create_manual(
+        scale=scale,
+        offset=offset,
+        input_stats_dict=stat
+    )
+
+def robomimic_abs_action_normalizer_from_stat(stat, rotation_transformer):
+    result = dict_apply_split(
+        stat, lambda x: {
+            'pos': x[...,:3],
+            'rot': x[...,3:6],
+            'gripper': x[...,6:]
+    })
+
+    def get_pos_param_info(stat, output_max=1, output_min=-1, range_eps=1e-7):
+        # -1, 1 normalization
+        input_max = stat['max']
+        input_min = stat['min']
+        input_range = input_max - input_min
+        ignore_dim = input_range < range_eps
+        input_range[ignore_dim] = output_max - output_min
+        scale = (output_max - output_min) / input_range
+        offset = output_min - scale * input_min
+        offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim]
+
+        return {'scale': scale, 'offset': offset}, stat
+
+    def get_rot_param_info(stat):
+        example = rotation_transformer.forward(stat['mean'])
+        scale = np.ones_like(example)
+        offset = np.zeros_like(example)
+        info = {
+            'max': np.ones_like(example),
+            'min': np.full_like(example, -1),
+            'mean': np.zeros_like(example),
+            'std': np.ones_like(example)
+        }
+        return {'scale': scale, 'offset': offset}, info
+    
+    def get_gripper_param_info(stat):
+        example = stat['max']
+        scale = np.ones_like(example)
+        offset = np.zeros_like(example)
+        info = {
+            'max': np.ones_like(example),
+            'min': np.full_like(example, -1),
+            'mean': np.zeros_like(example),
+            'std': np.ones_like(example)
+        }
+        return {'scale': scale, 'offset': offset}, info
+
+    pos_param, pos_info = get_pos_param_info(result['pos'])
+    rot_param, rot_info = get_rot_param_info(result['rot'])
+    gripper_param, gripper_info = get_gripper_param_info(result['gripper'])
+
+    param = dict_apply_reduce(
+        [pos_param, rot_param, gripper_param], 
+        lambda x: np.concatenate(x,axis=-1))
+    info = dict_apply_reduce(
+        [pos_info, rot_info, gripper_info], 
+        lambda x: np.concatenate(x,axis=-1))
+
+    return SingleFieldLinearNormalizer.create_manual(
+        scale=param['scale'],
+        offset=param['offset'],
+        input_stats_dict=info
+    )
+
+
+def robomimic_abs_action_only_normalizer_from_stat(stat):
+    result = dict_apply_split(
+        stat, lambda x: {
+            'pos': x[...,:3],
+            'other': x[...,3:]
+    })
+
+    def get_pos_param_info(stat, output_max=1, output_min=-1, range_eps=1e-7):
+        # -1, 1 normalization
+        input_max = stat['max']
+        input_min = stat['min']
+        input_range = input_max - input_min
+        ignore_dim = input_range < range_eps
+        input_range[ignore_dim] = output_max - output_min
+        scale = (output_max - output_min) / input_range
+        offset = output_min - scale * input_min
+        offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim]
+
+        return {'scale': scale, 'offset': offset}, stat
+
+    
+    def get_other_param_info(stat):
+        example = stat['max']
+        scale = np.ones_like(example)
+        offset = np.zeros_like(example)
+        info = {
+            'max': np.ones_like(example),
+            'min': np.full_like(example, -1),
+            'mean': np.zeros_like(example),
+            'std': np.ones_like(example)
+        }
+        return {'scale': scale, 'offset': offset}, info
+
+    pos_param, pos_info = get_pos_param_info(result['pos'])
+    other_param, other_info = get_other_param_info(result['other'])
+
+    param = dict_apply_reduce(
+        [pos_param, other_param], 
+        lambda x: np.concatenate(x,axis=-1))
+    info = dict_apply_reduce(
+        [pos_info, other_info], 
+        lambda x: np.concatenate(x,axis=-1))
+
+    return SingleFieldLinearNormalizer.create_manual(
+        scale=param['scale'],
+        offset=param['offset'],
+        input_stats_dict=info
+    )
+
+
+def robomimic_abs_action_only_dual_arm_normalizer_from_stat(stat):
+    Da = stat['max'].shape[-1]
+    Dah = Da // 2
+    result = dict_apply_split(
+        stat, lambda x: {
+            'pos0': x[...,:3],
+            'other0': x[...,3:Dah],
+            'pos1': x[...,Dah:Dah+3],
+            'other1': x[...,Dah+3:]
+    })
+
+    def get_pos_param_info(stat, output_max=1, output_min=-1, range_eps=1e-7):
+        # -1, 1 normalization
+        input_max = stat['max']
+        input_min = stat['min']
+        input_range = input_max - input_min
+        ignore_dim = input_range < range_eps
+        input_range[ignore_dim] = output_max - output_min
+        scale = (output_max - output_min) / input_range
+        offset = output_min - scale * input_min
+        offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim]
+
+        return {'scale': scale, 'offset': offset}, stat
+
+    
+    def get_other_param_info(stat):
+        example = stat['max']
+        scale = np.ones_like(example)
+        offset = np.zeros_like(example)
+        info = {
+            'max': np.ones_like(example),
+            'min': np.full_like(example, -1),
+            'mean': np.zeros_like(example),
+            'std': np.ones_like(example)
+        }
+        return {'scale': scale, 'offset': offset}, info
+
+    pos0_param, pos0_info = get_pos_param_info(result['pos0'])
+    pos1_param, pos1_info = get_pos_param_info(result['pos1'])
+    other0_param, other0_info = get_other_param_info(result['other0'])
+    other1_param, other1_info = get_other_param_info(result['other1'])
+
+    param = dict_apply_reduce(
+        [pos0_param, other0_param, pos1_param, other1_param], 
+        lambda x: np.concatenate(x,axis=-1))
+    info = dict_apply_reduce(
+        [pos0_info, other0_info, pos1_info, other1_info], 
+        lambda x: np.concatenate(x,axis=-1))
+
+    return SingleFieldLinearNormalizer.create_manual(
+        scale=param['scale'],
+        offset=param['offset'],
+        input_stats_dict=info
+    )
+
+
+def array_to_stats(arr: np.ndarray):
+    stat = {
+        'min': np.min(arr, axis=0),
+        'max': np.max(arr, axis=0),
+        'mean': np.mean(arr, axis=0),
+        'std': np.std(arr, axis=0)
+    }
+    return stat

common/pose_trajectory_interpolator.py

@@ -0,0 +1,208 @@
+from typing import Union
+import numbers
+import numpy as np
+import scipy.interpolate as si
+import scipy.spatial.transform as st
+
+def rotation_distance(a: st.Rotation, b: st.Rotation) -> float:
+    return (b * a.inv()).magnitude()
+
+def pose_distance(start_pose, end_pose):
+    start_pose = np.array(start_pose)
+    end_pose = np.array(end_pose)
+    start_pos = start_pose[:3]
+    end_pos = end_pose[:3]
+    start_rot = st.Rotation.from_rotvec(start_pose[3:])
+    end_rot = st.Rotation.from_rotvec(end_pose[3:])
+    pos_dist = np.linalg.norm(end_pos - start_pos)
+    rot_dist = rotation_distance(start_rot, end_rot)
+    return pos_dist, rot_dist
+
+class PoseTrajectoryInterpolator:
+    def __init__(self, times: np.ndarray, poses: np.ndarray):
+        assert len(times) >= 1
+        assert len(poses) == len(times)
+        if not isinstance(times, np.ndarray):
+            times = np.array(times)
+        if not isinstance(poses, np.ndarray):
+            poses = np.array(poses)
+
+        if len(times) == 1:
+            # special treatment for single step interpolation
+            self.single_step = True
+            self._times = times
+            self._poses = poses
+        else:
+            self.single_step = False
+            assert np.all(times[1:] >= times[:-1])
+
+            pos = poses[:,:3]
+            rot = st.Rotation.from_rotvec(poses[:,3:])
+
+            self.pos_interp = si.interp1d(times, pos, 
+                axis=0, assume_sorted=True)
+            self.rot_interp = st.Slerp(times, rot)
+    
+    @property
+    def times(self) -> np.ndarray:
+        if self.single_step:
+            return self._times
+        else:
+            return self.pos_interp.x
+    
+    @property
+    def poses(self) -> np.ndarray:
+        if self.single_step:
+            return self._poses
+        else:
+            n = len(self.times)
+            poses = np.zeros((n, 6))
+            poses[:,:3] = self.pos_interp.y
+            poses[:,3:] = self.rot_interp(self.times).as_rotvec()
+            return poses
+
+    def trim(self, 
+            start_t: float, end_t: float
+            ) -> "PoseTrajectoryInterpolator":
+        assert start_t <= end_t
+        times = self.times
+        should_keep = (start_t < times) & (times < end_t)
+        keep_times = times[should_keep]
+        all_times = np.concatenate([[start_t], keep_times, [end_t]])
+        # remove duplicates, Slerp requires strictly increasing x
+        all_times = np.unique(all_times)
+        # interpolate
+        all_poses = self(all_times)
+        return PoseTrajectoryInterpolator(times=all_times, poses=all_poses)
+    
+    def drive_to_waypoint(self, 
+            pose, time, curr_time,
+            max_pos_speed=np.inf, 
+            max_rot_speed=np.inf
+        ) -> "PoseTrajectoryInterpolator":
+        assert(max_pos_speed > 0)
+        assert(max_rot_speed > 0)
+        time = max(time, curr_time)
+        
+        curr_pose = self(curr_time)
+        pos_dist, rot_dist = pose_distance(curr_pose, pose)
+        pos_min_duration = pos_dist / max_pos_speed
+        rot_min_duration = rot_dist / max_rot_speed
+        duration = time - curr_time
+        duration = max(duration, max(pos_min_duration, rot_min_duration))
+        assert duration >= 0
+        last_waypoint_time = curr_time + duration
+
+        # insert new pose
+        trimmed_interp = self.trim(curr_time, curr_time)
+        times = np.append(trimmed_interp.times, [last_waypoint_time], axis=0)
+        poses = np.append(trimmed_interp.poses, [pose], axis=0)
+
+        # create new interpolator
+        final_interp = PoseTrajectoryInterpolator(times, poses)
+        return final_interp
+
+    def schedule_waypoint(self,
+            pose, time, 
+            max_pos_speed=np.inf, 
+            max_rot_speed=np.inf,
+            curr_time=None,
+            last_waypoint_time=None
+        ) -> "PoseTrajectoryInterpolator":
+        assert(max_pos_speed > 0)
+        assert(max_rot_speed > 0)
+        if last_waypoint_time is not None:
+            assert curr_time is not None
+
+        # trim current interpolator to between curr_time and last_waypoint_time
+        start_time = self.times[0]
+        end_time = self.times[-1]
+        assert start_time <= end_time
+
+        if curr_time is not None:
+            if time <= curr_time:
+                # if insert time is earlier than current time
+                # no effect should be done to the interpolator
+                return self
+            # now, curr_time < time
+            start_time = max(curr_time, start_time)
+
+            if last_waypoint_time is not None:
+                # if last_waypoint_time is earlier than start_time
+                # use start_time
+                if time <= last_waypoint_time:
+                    end_time = curr_time
+                else:
+                    end_time = max(last_waypoint_time, curr_time)
+            else:
+                end_time = curr_time
+
+        end_time = min(end_time, time)
+        start_time = min(start_time, end_time)
+        # end time should be the latest of all times except time
+        # after this we can assume order (proven by zhenjia, due to the 2 min operations)
+
+        # Constraints:
+        # start_time <= end_time <= time (proven by zhenjia)
+        # curr_time <= start_time (proven by zhenjia)
+        # curr_time <= time (proven by zhenjia)
+        
+        # time can't change
+        # last_waypoint_time can't change
+        # curr_time can't change
+        assert start_time <= end_time
+        assert end_time <= time
+        if last_waypoint_time is not None:
+            if time <= last_waypoint_time:
+                assert end_time == curr_time
+            else:
+                assert end_time == max(last_waypoint_time, curr_time)
+
+        if curr_time is not None:
+            assert curr_time <= start_time
+            assert curr_time <= time
+
+        trimmed_interp = self.trim(start_time, end_time)
+        # after this, all waypoints in trimmed_interp is within start_time and end_time
+        # and is earlier than time
+
+        # determine speed
+        duration = time - end_time
+        end_pose = trimmed_interp(end_time)
+        pos_dist, rot_dist = pose_distance(pose, end_pose)
+        pos_min_duration = pos_dist / max_pos_speed
+        rot_min_duration = rot_dist / max_rot_speed
+        duration = max(duration, max(pos_min_duration, rot_min_duration))
+        assert duration >= 0
+        last_waypoint_time = end_time + duration
+
+        # insert new pose
+        times = np.append(trimmed_interp.times, [last_waypoint_time], axis=0)
+        poses = np.append(trimmed_interp.poses, [pose], axis=0)
+
+        # create new interpolator
+        final_interp = PoseTrajectoryInterpolator(times, poses)
+        return final_interp
+
+
+    def __call__(self, t: Union[numbers.Number, np.ndarray]) -> np.ndarray:
+        is_single = False
+        if isinstance(t, numbers.Number):
+            is_single = True
+            t = np.array([t])
+        
+        pose = np.zeros((len(t), 6))
+        if self.single_step:
+            pose[:] = self._poses[0]
+        else:
+            start_time = self.times[0]
+            end_time = self.times[-1]
+            t = np.clip(t, start_time, end_time)
+
+            pose = np.zeros((len(t), 6))
+            pose[:,:3] = self.pos_interp(t)
+            pose[:,3:] = self.rot_interp(t).as_rotvec()
+
+        if is_single:
+            pose = pose[0]
+        return pose

common/precise_sleep.py

@@ -0,0 +1,25 @@
+import time
+
+def precise_sleep(dt: float, slack_time: float=0.001, time_func=time.monotonic):
+    """
+    Use hybrid of time.sleep and spinning to minimize jitter.
+    Sleep dt - slack_time seconds first, then spin for the rest.
+    """
+    t_start = time_func()
+    if dt > slack_time:
+        time.sleep(dt - slack_time)
+    t_end = t_start + dt
+    while time_func() < t_end:
+        pass
+    return
+
+def precise_wait(t_end: float, slack_time: float=0.001, time_func=time.monotonic):
+    t_start = time_func()
+    t_wait = t_end - t_start
+    if t_wait > 0:
+        t_sleep = t_wait - slack_time
+        if t_sleep > 0:
+            time.sleep(t_sleep)
+        while time_func() < t_end:
+            pass
+    return

common/pref_replay_buffer.py

@@ -0,0 +1,365 @@
+from typing import Union, Dict, Optional
+import os
+import math
+import numbers
+import zarr
+import numcodecs
+import numpy as np
+from functools import cached_property
+
+def check_chunks_compatible(chunks: tuple, shape: tuple):
+    assert len(shape) == len(chunks)
+    for c in chunks:
+        assert isinstance(c, numbers.Integral)
+        assert c > 0
+
+def rechunk_recompress_array(group, name, 
+        chunks=None, chunk_length=None,
+        compressor=None, tmp_key='_temp'):
+    old_arr = group[name]
+    if chunks is None:
+        if chunk_length is not None:
+            chunks = (chunk_length,) + old_arr.chunks[1:]
+        else:
+            chunks = old_arr.chunks
+    check_chunks_compatible(chunks, old_arr.shape)
+    
+    if compressor is None:
+        compressor = old_arr.compressor
+    
+    if (chunks == old_arr.chunks) and (compressor == old_arr.compressor):
+        # no change
+        return old_arr
+
+    # rechunk recompress
+    group.move(name, tmp_key)
+    old_arr = group[tmp_key]
+    n_copied, n_skipped, n_bytes_copied = zarr.copy(
+        source=old_arr,
+        dest=group,
+        name=name,
+        chunks=chunks,
+        compressor=compressor,
+    )
+    del group[tmp_key]
+    arr = group[name]
+    return arr
+
+def get_optimal_chunks(shape, dtype, 
+        target_chunk_bytes=2e6, 
+        max_chunk_length=None):
+    """
+    Common shapes
+    T,D
+    T,N,D
+    T,H,W,C
+    T,N,H,W,C
+    """
+    itemsize = np.dtype(dtype).itemsize
+    # reversed
+    rshape = list(shape[::-1])
+    if max_chunk_length is not None:
+        rshape[-1] = int(max_chunk_length)
+    split_idx = len(shape)-1
+    for i in range(len(shape)-1):
+        this_chunk_bytes = itemsize * np.prod(rshape[:i])
+        next_chunk_bytes = itemsize * np.prod(rshape[:i+1])
+        if this_chunk_bytes <= target_chunk_bytes \
+            and next_chunk_bytes > target_chunk_bytes:
+            split_idx = i
+
+    rchunks = rshape[:split_idx]
+    item_chunk_bytes = itemsize * np.prod(rshape[:split_idx])
+    this_max_chunk_length = rshape[split_idx]
+    next_chunk_length = min(this_max_chunk_length, math.ceil(
+            target_chunk_bytes / item_chunk_bytes))
+    rchunks.append(next_chunk_length)
+    len_diff = len(shape) - len(rchunks)
+    rchunks.extend([1] * len_diff)
+    chunks = tuple(rchunks[::-1])
+    # print(np.prod(chunks) * itemsize / target_chunk_bytes)
+    return chunks
+
+
+class PrefReplayBuffer:
+    """
+    Zarr-based temporal data structure specifically for preference dataset.
+    Stores pairs of trajectories (observations, actions) along with votes.
+    """
+    
+    def __init__(self, root: Union[zarr.Group, Dict[str, dict]]):
+        """
+        Initialize the preference replay buffer. Use class methods to create or load buffers.
+        """
+        assert 'data' in root
+        assert 'meta' in root
+        for key, value in root['data'].items():
+            assert value.shape[0] == root['meta']['votes'].shape[0]
+        self.root = root
+
+    # ============= create constructors ===============
+    @classmethod
+    def create_empty_zarr(cls, storage=None, root=None):
+        if root is None:
+            if storage is None:
+                storage = zarr.MemoryStore()
+            root = zarr.group(store=storage)
+        data = root.require_group('data', overwrite=False)
+        meta = root.require_group('meta', overwrite=False)
+        
+        if 'votes' not in meta:
+            votes = meta.zeros('votes', shape=(0,), dtype=np.float32, compressor=None, overwrite=False)
+        if 'votes_2' not in meta:
+            votes_2 = meta.zeros('votes_2', shape=(0,), dtype=np.float32, compressor=None, overwrite=False)
+        
+        return cls(root=root)
+
+    @classmethod
+    def create_empty_numpy(cls):
+        root = {
+            'data': dict(),
+            'meta': {
+                'episode_ends': np.zeros((0,), dtype=np.int64),
+                'votes': np.zeros((0,), dtype=np.float32), 
+                'votes_2': np.zeros((0,), dtype=np.float32),
+                'length': np.zeros((0,), dtype=np.int64),
+                'length_2': np.zeros((0,), dtype=np.int64),
+                'beta_priori': np.zeros((0,), dtype=np.float32),
+                'beta_priori_2': np.zeros((0,), dtype=np.float32),
+            }
+        }
+        return cls(root=root)
+
+    @classmethod
+    def create_from_group(cls, group, **kwargs):
+        if 'data' not in group:
+            # create from scratch
+            buffer = cls.create_empty_zarr(root=group, **kwargs)
+        else:
+            # already exists
+            buffer = cls(root=group, **kwargs)
+        return buffer
+
+    @classmethod
+    def create_from_path(cls, zarr_path, mode='r', **kwargs):
+        """
+        Open a Zarr file from disk for large datasets that cannot fit in memory.
+        """
+        group = zarr.open(zarr_path, mode=mode)
+        return cls.create_from_group(group=group, **kwargs)
+
+    # ============= Add episodes ===============
+    def add_pref_episode(self, data: Dict[str, np.ndarray], 
+                         meta_data: Optional[Dict[str, Union[np.ndarray, int]]] = None,
+                         chunks: Optional[Dict[str, tuple]] = dict(),
+                         compressors: Union[str, numcodecs.abc.Codec, dict] = dict()):
+        """
+        Add a pair of episodes (obs/action for each trajectory) along with metadata (votes).
+        """
+        assert 'obs' in data and 'obs_2' in data, "obs and obs_2 keys are required"
+        assert 'action' in data and 'action_2' in data, "action and action_2 keys are required"
+        
+        is_zarr = isinstance(self.root, zarr.Group)
+        curr_len = len(self.root['meta']['votes'])
+        episode_length = len(data['obs'])
+        new_len = curr_len + 1
+
+        # Add trajectory 1
+        for key in ['obs', 'action']:
+            # Create the new shape to accommodate all time steps
+            value = data[key]
+            new_shape = (new_len,) + (episode_length,) + data[key].shape[1:]  # This will set (new_len, T, dim)
+
+            if key not in self.root['data']:
+                # Create a new array if it doesn't exist
+                if is_zarr:
+                    cks = self._resolve_array_chunks(chunks, key, data[key])
+                    cpr = self._resolve_array_compressor(compressors, key, data[key])
+                    arr = self.root['data'].zeros(name=key, shape=new_shape, chunks=cks, dtype=data[key].dtype, compressor=cpr)
+                else:
+                    arr = np.zeros(new_shape, dtype=data[key].dtype)
+                    self.root['data'][key] = arr
+            else:
+                arr = self.root['data'][key]
+                if is_zarr:
+                    arr.resize(new_shape)
+                else:
+                    arr.resize(new_shape, refcheck=False)
+
+            # Store the full sequence, adjusting the shape to match the time steps in data[key]
+            arr[new_len-1, -value.shape[0]:, :] = value # Now this assumes data[key] has shape (T, dim)
+
+
+        # Add trajectory 2 (obs_2, action_2)
+        for key in ['obs_2', 'action_2']:
+            value = data[key]
+            # Create the new shape to accommodate all time steps
+            new_shape = (new_len,) + (episode_length,) + data[key].shape[1:]  # This will set (new_len, T, dim)
+
+            if key not in self.root['data']:
+                # Create a new array if it doesn't exist
+                if is_zarr:
+                    cks = self._resolve_array_chunks(chunks, key, data[key])
+                    cpr = self._resolve_array_compressor(compressors, key, data[key])
+                    arr = self.root['data'].zeros(name = key, shape=new_shape, chunks=cks, dtype=data[key].dtype, compressor=cpr)
+                else:
+                    arr = np.zeros(new_shape, dtype=data[key].dtype)
+                    self.root['data'][key] = arr
+            else:
+                arr = self.root['data'][key]
+                if is_zarr:
+                    arr.resize(new_shape)
+                else:
+                    arr.resize(new_shape, refcheck=False)
+
+            # Store the full sequence, adjusting the shape to match the time steps in data[key]
+            arr[new_len-1, -value.shape[0]:, :] = value  # Now this assumes data[key] has shape (T, dim)
+
+        # Add votes to meta
+        if meta_data:
+            for key in ['votes', 'votes_2']:
+                new_shape = (new_len,) + (1,)
+                if key not in self.root['meta']:
+                    if is_zarr:
+                        self.root['meta'].zeros(name=key, shape=new_shape, chunks=new_shape, dtype=np.float32)
+                    else:
+                        self.root['meta'][key] = np.zeros(new_shape, dtype=np.float32)
+                arr = self.root['meta'][key]
+                if is_zarr:
+                    arr.resize(new_shape)
+                else:
+                    arr.resize(new_shape, refcheck=False)
+                arr[new_len-1] = meta_data[key]
+
+            for key in ['length', 'length_2']:
+                new_shape = (new_len,) + (1,)
+                if key not in self.root['meta']:
+                    if is_zarr:
+                        self.root['meta'].zeros(name=key, shape=new_shape, chunks=new_shape, dtype=np.float32)
+                    else:
+                        self.root['meta'][key] = np.zeros(new_shape, dtype=np.float32)
+                arr = self.root['meta'][key]
+                if is_zarr:
+                    arr.resize(new_shape)
+                else:
+                    arr.resize(new_shape, refcheck=False)
+                arr[new_len-1] = meta_data[key]
+
+            for key in ['beta_priori', 'beta_priori_2']:
+                new_shape = (new_len,) + (2,)
+                if key not in self.root['meta']:
+                    if is_zarr:
+                        self.root['meta'].zeros(name=key, shape=new_shape, chunks=new_shape, dtype=np.float32)
+                    else: 
+                        self.root['meta'][key] = np.zeros(new_shape, dtype=np.float32)
+                arr = self.root['meta'][key]
+                if is_zarr:
+                    arr.resize(new_shape)
+                else:
+                    arr.resize(new_shape, refcheck=False)
+                arr[new_len-1] = meta_data[key]
+
+    # ============= Get episodes ===============
+    def get_pref_episode(self, idx: int, copy: bool = False):
+        """
+        Get a pair of episodes by index, including observation and action sequences for both trajectories.
+        """
+        if copy:
+            return {
+                'obs': self.root['data']['obs'][idx].copy(),
+                'action': self.root['data']['action'][idx].copy(),
+                'obs_2': self.root['data']['obs_2'][idx].copy(),
+                'action_2': self.root['data']['action_2'][idx].copy(),
+                'votes': self.root['meta']['votes'][idx].copy(),
+                'votes_2': self.root['meta']['votes_2'][idx].copy(),
+                'length': self.root['meta']['length'][idx].copy(),
+                'length_2': self.root['meta']['length_2'][idx].copy(),
+                'beta_priori': self.root['meta']['beta_priori'][idx].copy(),
+                'beta_priori_2': self.root['meta']['beta_priori_2'][idx].copy(),
+            }
+        else:
+            return {
+                'obs': self.root['data']['obs'][idx],
+                'action': self.root['data']['action'][idx],
+                'obs_2': self.root['data']['obs_2'][idx],
+                'action_2': self.root['data']['action_2'][idx],
+                'votes': self.root['meta']['votes'][idx],
+                'votes_2': self.root['meta']['votes_2'][idx],
+                'length': self.root['meta']['length'][idx],
+                'length_2': self.root['meta']['length_2'][idx],
+                'beta_priori': self.root['meta']['beta_priori'][idx],
+                'beta_priori_2': self.root['meta']['beta_priori_2'][idx],
+            }
+
+    def get_episode_slice(self, idx):
+        """
+        Get the slice range for an episode based on the index for slicing observation and action arrays.
+        """
+        return slice(idx, idx + 1)
+
+    # ============= Save methods ===============
+    def save_to_store(self, store, chunks: Optional[Dict[str, tuple]] = dict(),
+                      compressors: Union[str, numcodecs.abc.Codec, dict] = dict(),
+                      if_exists='replace', **kwargs):
+        """
+        Save the replay buffer to the given store with optional chunking and compression.
+        """
+        root = zarr.group(store=store)
+        # Save data and meta with chunking and compression if provided
+        for key, value in self.root['data'].items():
+            cks = self._resolve_array_chunks(chunks, key, value)
+            cpr = self._resolve_array_compressor(compressors, key, value)
+            arr = self.root['data'][key]
+            zarr.copy(arr, root['data'], name=key, chunks=cks, compressor=cpr)
+        
+        # Save meta
+        for key, value in self.root['meta'].items():
+            zarr.copy(value, root['meta'], name=key)
+    
+    def save_to_path(self, zarr_path, chunks: Optional[Dict[str, tuple]] = dict(),
+                     compressors: Union[str, numcodecs.abc.Codec, dict] = dict(), 
+                     if_exists='replace', **kwargs):
+        """
+        Save the replay buffer to a path.
+        """
+        store = zarr.DirectoryStore(zarr_path)
+        return self.save_to_store(store, chunks=chunks, compressors=compressors, if_exists=if_exists, **kwargs)
+
+    # ============= Helper methods ===============
+    @classmethod
+    def _resolve_array_chunks(cls,
+            chunks: Union[dict, tuple], key, array):
+        cks = None
+        if isinstance(chunks, dict):
+            if key in chunks:
+                cks = chunks[key]
+            elif isinstance(array, zarr.Array):
+                cks = array.chunks
+        elif isinstance(chunks, tuple):
+            cks = chunks
+        else:
+            raise TypeError(f"Unsupported chunks type {type(chunks)}")
+        # backup default
+        if cks is None:
+            cks = get_optimal_chunks(shape=array.shape, dtype=array.dtype)
+        # check
+        check_chunks_compatible(chunks=cks, shape=array.shape)
+        return cks
+
+    @classmethod
+    def _resolve_array_compressor(cls, compressors: Union[dict, str, numcodecs.abc.Codec], key, array):
+        # Resolves the compressor for the array
+        return compressors.get(key, array.compressor if isinstance(array, zarr.Array) else numcodecs.Blosc())
+
+    @property
+    def n_steps(self):
+        return len(self.root['meta']['votes'])
+
+    @property
+    def data(self):
+        return self.root['data']
+
+    @property
+    def meta(self):
+        return self.root['meta']
+

common/pref_sampler.py

@@ -0,0 +1,74 @@
+from typing import Optional, Dict
+import numpy as np
+from diffusion_policy.common.pref_replay_buffer import PrefReplayBuffer
+import torch
+
+
+def get_val_mask(n_episodes, val_ratio, seed=0):
+    val_mask = np.zeros(n_episodes, dtype=bool)
+    if val_ratio <= 0:
+        return val_mask
+
+    # have at least 1 episode for validation, and at least 1 episode for train
+    n_val = min(max(1, round(n_episodes * val_ratio)), n_episodes-1)
+    rng = np.random.default_rng(seed=seed)
+    val_idxs = rng.choice(n_episodes, size=n_val, replace=False)
+    val_mask[val_idxs] = True
+    return val_mask
+
+
+class PrefSequenceSampler:
+    def __init__(self,
+                 replay_buffer: PrefReplayBuffer,
+                 sequence_length: int,
+                 episode_mask: Optional[np.ndarray]=None,
+                 keys: Optional[Dict[str, int]] = None,
+                 ):
+        """
+        Initializes a sampler for the preference replay buffer.
+        
+        Parameters:
+        - replay_buffer: PrefReplayBuffer instance from which to sample data.
+        - sequence_length: The length of sequences to sample.
+        - pad_before, pad_after: Padding before and after sequences (optional).
+        - keys: Optional dictionary to specify specific keys and limits on how much data to load.
+        - episode_mask: Mask indicating valid episodes for sampling.
+        """
+        super().__init__()
+        assert sequence_length >= 1
+
+        if keys is None:
+            keys = list(replay_buffer.data.keys())
+
+        # Store generated indices
+        self.keys = keys
+        self.sequence_length = sequence_length
+        self.replay_buffer = replay_buffer
+        self.episode_mask = episode_mask
+
+    def __len__(self):
+        
+        return np.sum(self.episode_mask)
+
+    def sample_sequence(self, idx: int) -> Dict[str, np.ndarray]:
+        """
+        Samples the sequence of data based on the provided index (idx).
+        
+        Parameters:
+        - idx: The index from which to sample an episode sequence.
+        
+        Returns:
+        - A dictionary containing the sampled data for the specified keys and votes.
+        """
+        indices = np.where(self.episode_mask)[0]
+        result = self.replay_buffer.get_pref_episode(indices[idx])
+
+        for key in result:
+            value = result[key]
+            if isinstance(value, np.ndarray):
+                result[key] = torch.from_numpy(value)
+            elif isinstance(value, (np.float32, np.float64, float, int)):
+                result[key] = torch.tensor(value, dtype=torch.float32)
+            else:
+                raise TypeError(f"Unsupported type {type(value)} for key '{key}'")
+        return result

common/prior_utils_confidence.py

@@ -0,0 +1,723 @@
+import collections
+import numpy as np
+# import gym
+import torch.utils
+from tqdm import trange
+import torch
+import torch.nn as nn
+import math
+# import d4rl
+from tqdm import tqdm
+import torch.nn.functional as F
+from pathlib import Path
+import copy
+from torch.optim.lr_scheduler import SequentialLR, LinearLR, CosineAnnealingLR
+from diffusion_policy.policy.diffusion_transformer_hybrid_image_policy import DiffusionTransformerHybridImagePolicy
+from diffusion_policy.common.pytorch_util import dict_apply
+
+import time
+import logging
+from datetime import datetime
+
+Batch = collections.namedtuple(
+    'Batch',
+    ['observations', 'actions', 'rewards', 'masks', 'next_observations'])
+
+
+def gen_net(in_size=1, out_size=1, H=128, n_layers=3, activation='tanh'):
+    net = []
+    for i in range(n_layers):
+        net.append(nn.Linear(in_size, H))
+        net.append(nn.LeakyReLU())
+        in_size = H
+    net.append(nn.Linear(in_size, out_size))
+    if activation == 'tanh':
+        net.append(nn.Tanh())
+    elif activation == 'sig':
+        net.append(nn.Sigmoid())
+    else:
+        pass
+
+    return net
+
+def get_discrete_traj(trajectory, dim=2):
+    discrete_traj = np.round(trajectory[:, :, :dim].astype(float)).astype(int)
+
+    unique_states_per_trajectory = []
+    for trajectory in discrete_traj:
+        # Find unique states in the current trajectory
+        unique_states = np.unique(trajectory, axis=0)
+        # Convert the trajectory to a hashable type
+        unique_states_per_trajectory.append(tuple(map(tuple, unique_states)))
+    return unique_states_per_trajectory
+
+def init_trajectory_dict(discrete_trajectory):
+
+    trajectory_dict = {}
+    for trajectory in discrete_trajectory:
+        if trajectory in trajectory_dict:
+            continue
+        else:
+            trajectory_dict[trajectory] = np.array([1,1])
+    return trajectory_dict
+
+def get_trajectory_dict_from_pair(traj_alpha_beta_dict, discrete_obs_1, discrete_obs_2, single_labels):
+    for i in range(len(single_labels)):
+        if single_labels[i] == 1:
+            traj_alpha_beta_dict[discrete_obs_1[i]][0] += 1
+            traj_alpha_beta_dict[discrete_obs_2[i]][1] += 1
+        elif single_labels[i] == -1:
+            traj_alpha_beta_dict[discrete_obs_1[i]][1] += 1
+            traj_alpha_beta_dict[discrete_obs_2[i]][0] += 1
+        # TODO how to treat not compaied ones?
+        elif single_labels[i] == 0:
+            traj_alpha_beta_dict[discrete_obs_1[i]][0] += 1
+            traj_alpha_beta_dict[discrete_obs_1[i]][1] += 1
+            traj_alpha_beta_dict[discrete_obs_2[i]][0] += 1
+            traj_alpha_beta_dict[discrete_obs_2[i]][1] += 1
+    return traj_alpha_beta_dict
+
+def to_torch(x, dtype=torch.float32):
+    if isinstance(x, np.ndarray):
+        return torch.from_numpy(x).to(dtype)
+
+def index_batch(batch, indices):
+    indexed = {}
+    for key in batch.keys():
+        indexed[key] = batch[key][indices, ...]
+    return indexed
+
+class Dataset(object):
+    def __init__(self, observations: np.ndarray, actions: np.ndarray,
+                 rewards: np.ndarray, masks: np.ndarray,
+                 dones_float: np.ndarray, next_observations: np.ndarray,
+                 size: int):
+        self.observations = observations
+        self.actions = actions
+        self.rewards = rewards
+        self.masks = masks
+        self.dones_float = dones_float
+        self.next_observations = next_observations
+        self.size = size
+
+    def sample(self, batch_size: int) -> Batch:
+        indx = np.random.randint(self.size, size=batch_size)
+        return Batch(observations=self.observations[indx],
+                     actions=self.actions[indx],
+                     rewards=self.rewards[indx],
+                     masks=self.masks[indx],
+                     next_observations=self.next_observations[indx])
+
+class NormalComparisonModel(nn.Module):
+    def __init__(self, input_dim, dense_units, dropout_rate, device):
+        super(NormalComparisonModel, self).__init__()
+        
+        # Fully connected layers for classification
+        self.fc1 = nn.Linear(4 * input_dim, dense_units)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.fc2 = nn.Linear(dense_units, dense_units // 2)
+        self.output = nn.Linear(dense_units // 2, 1)
+        
+        # Device setup
+        self.device = device
+        self.to(device)
+
+    def forward(self, f1, f2):
+        f1, f2 = f1.to(self.device), f2.to(self.device)
+        
+        # Shape of f1: (N1, L1, D), f2: (N2, L2, D)
+        N1, L1, D = f1.shape
+        N2, L2, _ = f2.shape
+
+        # Expand dimensions for pairwise comparison
+        f1_expanded = f1.unsqueeze(1).expand(N1, N2, L1, D)  # Shape: (N1, N2, L1, D)
+        f2_expanded = f2.unsqueeze(0).expand(N1, N2, L2, D)  # Shape: (N1, N2, L2, D)
+        
+        # Flatten for processing through Transformer
+        f1_flat = f1_expanded.reshape(-1, L1, D)  # Shape: (N1*N2, L1, D)
+        f2_flat = f2_expanded.reshape(-1, L2, D)  # Shape: (N1*N2, L2, D)
+
+        # Sequence pooling: Reduce sequence dimension
+        f1_pooled = f1_flat.mean(dim=1)
+        f2_pooled = f2_flat.mean(dim=1)
+
+        # Pairwise comparison (concatenate, subtract)
+        combined_features = torch.cat([
+            f1_pooled,  # Reduced f1
+            f2_pooled,  # Reduced f2
+            f1_pooled - f2_pooled,  # Difference
+        ], dim=-1)  # Shape: (N1*N2, 3*D)
+
+        # Fully connected layers
+        x = F.gelu(self.fc1(combined_features))  # Shape: (N1*N2, dense_units)
+        x = self.dropout(x)
+        x = F.gelu(self.fc2(x))  # Shape: (N1*N2, dense_units // 2)
+        x = self.output(x)  # Shape: (N1*N2, 1)
+        
+        # Sigmoid activation for [0, 1] output
+        output = torch.sigmoid(x).squeeze(-1)  # Shape: (N1*N2)
+        output = output.view(N1, N2)  # Reshape to (N1, N2)
+        return output
+
+class AttentionComparisonModel(nn.Module):
+    def __init__(self, input_dim, dropout_rate, nhead, device):
+        super(AttentionComparisonModel, self).__init__()
+
+        self.attention = nn.MultiheadAttention(embed_dim=input_dim, num_heads=nhead, batch_first=True)
+        self.query_attn = nn.Parameter(torch.randn(1, input_dim))  # Learnable query for attention
+
+        self.fc1 = nn.Linear(input_dim, input_dim // 2)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.output = nn.Linear(input_dim // 2, 1)
+        self.device = device
+        self.to(device)
+
+    def attention_pool(self, x, query):
+        # x: (batch_size, seq_len, input_dim)
+        # query: (1, input_dim)
+        attn_weights = torch.softmax(torch.matmul(x, query.T), dim=1)  # Compute attention weights
+        weighted_sum = torch.sum(attn_weights * x, dim=1)  # Weighted sum along sequence dimension
+        return weighted_sum
+
+    def forward(self, f1, f2):
+        f1, f2 = f1.to(self.device), f2.to(self.device)
+
+        N1, L1, D = f1.shape
+        N2, L2, _ = f2.shape
+
+        f1_flat = f1.unsqueeze(1).expand(N1, N2, L1, D).reshape(-1, L1, D)
+        f2_flat = f2.unsqueeze(0).expand(N1, N2, L2, D).reshape(-1, L2, D)
+
+        attn_out, _ = self.attention(f1_flat, f2_flat, f2_flat)
+
+        attn_reduced = self.attention_pool(attn_out, self.query_attn)
+
+        x = F.gelu(self.fc1(attn_reduced))
+        x = self.dropout(x)
+        x = self.output(x)
+        output = torch.sigmoid(x)
+        output = output.view(N1, N2)
+        return output
+
+    def one_to_one_forward(self, f1, f2):
+        f1, f2 = f1.to(self.device), f2.to(self.device)
+
+        attn_out, _ = self.attention(f1, f2, f2)
+
+        # Apply self-attention weighted compression
+        attn_reduced = self.attention_pool(attn_out, self.query_attn)
+
+        x = F.gelu(self.fc1(attn_reduced))
+        x = self.dropout(x)
+        x = self.output(x)
+        output = torch.sigmoid(x)
+        return output
+
+
+class SinusoidalPosEmb(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = x[:, None] * emb[None, :]
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+class TransformerEncModel(nn.Module):
+    def __init__(self, data_dim, embedding_dim, nhead, num_encoder_layers, device):
+        super(TransformerEncModel, self).__init__()
+        self.device = device
+        self.embedding = nn.Linear(data_dim, embedding_dim)
+        # self.pos_encoder = self.create_positional_encoding(seq_length, embedding_dim)
+        self.pos_emb = SinusoidalPosEmb(embedding_dim)
+        encoder_layers = nn.TransformerEncoderLayer(d_model=embedding_dim, nhead=nhead, batch_first=True)
+        self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_layers=num_encoder_layers)
+
+    def forward(self, x):
+        traj_len = x.shape[1]
+        x = self.embedding(x)  # Map input to embedding dimension
+        pos = self.pos_emb(
+            torch.arange(traj_len, device=self.device))[None,]
+        x += pos  # Add positional encoding
+        output = self.transformer_encoder(x)
+        return output
+
+class CausalTransformerBetaModel(nn.Module):
+    def __init__(self, data_dim, embedding_dim, nhead, num_encoder_layers, output_dim, device):
+        super(CausalTransformerBetaModel, self).__init__()
+        self.device = device
+        self.embedding = nn.Linear(data_dim, embedding_dim)
+        # self.pos_encoder = self.create_positional_encoding(seq_length, embedding_dim)
+        self.pos_emb = SinusoidalPosEmb(embedding_dim)
+        encoder_layers = nn.TransformerEncoderLayer(d_model=embedding_dim, nhead=nhead, batch_first=True)
+        self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_layers=num_encoder_layers)
+        self.output_layer = nn.Linear(embedding_dim, output_dim)
+        self.softplus = nn.Softplus()
+
+    def generate_square_subsequent_mask(self, sz):
+        mask = torch.triu(torch.ones(sz, sz), diagonal=1)
+        return mask.masked_fill(mask==1, float('-inf'))
+
+    def forward(self, x):
+        traj_len = x.shape[1]
+        x = self.embedding(x)  # Map input to embedding dimension
+        pos = self.pos_emb(
+            torch.arange(traj_len, device=self.device))[None,]
+        x += pos  # Add positional encoding
+        mask = self.generate_square_subsequent_mask(traj_len).to(x.device)
+        output = self.transformer_encoder(x, mask=mask)
+        output = self.output_layer(output)
+        # TODO mean or last one
+        # output = output.mean(dim=1)  # Aggregate across the sequence dimension
+        output = output[:, -1, :]  # Take the output from the last timestep
+        output = self.softplus(output)
+        return output
+
+class BetaNetwork(nn.Module):
+    def __init__(self, data, device=torch.device('cuda'), data_size = 500):
+        super(BetaNetwork, self).__init__()
+
+        act_data = np.concatenate((data['action'], data['action_2']), axis=0)
+        obs_data = np.concatenate((data['obs'], data['obs_2']), axis=0)
+        votes_data = np.concatenate((data['votes'], data['votes_2']), axis=0)
+
+        if data_size <= obs_data.shape[0]:
+            indices = np.random.randint(0, obs_data.shape[0], size=data_size)
+            obs_data = obs_data[indices, ...]
+            act_data = act_data[indices, ...]
+            votes_data = votes_data[indices, ...]
+
+        act_data = torch.from_numpy(act_data).float().to(device)
+        obs_data = torch.from_numpy(obs_data).float().to(device)
+        self.votes_data = torch.from_numpy(votes_data).to(device)
+        self.lr = None
+        self.device = device
+        self.data = torch.concat((obs_data, act_data), dim=-1)
+
+        class BetaModel(nn.Module):
+            def __init__(self, obs_data, act_data, device=torch.device('cuda')):
+                super(BetaModel, self).__init__()
+
+                self.enc_model = TransformerEncModel(
+                    data_dim = act_data.shape[-1] + obs_data.shape[-1],
+                    embedding_dim = 256,
+                    nhead = 4,
+                    num_encoder_layers = 2,
+                    device = device
+                ).to(device)
+
+                self.comp_model = AttentionComparisonModel(
+                    input_dim = 256, 
+                    dropout_rate = 0.3,
+                    nhead = 16,
+                    device = device
+                ).to(device)
+
+                self.data = torch.concat((obs_data, act_data), dim=-1)
+
+            def forward(self, x):
+                batch_f = self.enc_model(x)
+                all_data_f = self.enc_model(self.data)
+                bias = all_data_f.mean()
+                std = all_data_f.std()
+                all_data_f = (all_data_f - bias) / std
+                batch_f = (batch_f - bias) / std
+                output = self.comp_model(batch_f, all_data_f)
+                return output
+            
+            def one_to_one_forward(self, x, y):
+                x_f = self.enc_model(x)
+                y_f = self.enc_model(y)
+                all_data_f = self.enc_model(self.data)
+                bias = all_data_f.mean()
+                std = all_data_f.std()
+                x_f = (x_f - bias) / std
+                y_f = (y_f - bias) / std
+                output = self.comp_model.one_to_one_forward(x_f, y_f)
+                return output
+
+       
+
+        self.model = BetaModel(obs_data, act_data, device)
+        self.opt = None #, weight_decay=1.0e-4
+        self.scheduler = None #torch.optim.lr_scheduler.ReduceLROnPlateau(self.opt, mode='min', patience=100, verbose=True)
+        self.ref_model: BetaModel
+
+    def get_alpha_beta(self, x):
+        batch_comp = self.model(x).detach()
+        alpha = torch.sum(batch_comp, dim=-1)
+        beta = torch.sum(1 - batch_comp, dim=-1)
+        # alpha = torch.sum(torch.where(batch_comp > 0.5, torch.tensor(1.0, device=self.device), torch.tensor(0.0, device=self.device)), dim=-1)
+        # beta = torch.sum(torch.where(batch_comp < 0.5, torch.tensor(1.0, device=self.device), torch.tensor(0.0, device=self.device)), dim=-1)
+
+        return alpha.detach(), beta.detach()
+
+    def fit_data(self, save_dir=None, load_dir=None, num_epochs=1, warm_up_epochs=0, batch_size=1, lr=1.0e-5):
+        if load_dir is None:
+            interval = math.ceil(self.data.shape[0] / batch_size)
+            total_steps = num_epochs * interval
+            warm_up_steps = warm_up_epochs * interval
+            main_steps = total_steps - warm_up_steps
+
+            # Learning rate schedulers
+            self.lr = lr
+            self.opt = torch.optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=1e-5)
+            warm_up_scheduler = LinearLR(self.opt, start_factor=1e-8, end_factor=1.0, total_iters=warm_up_steps)
+            cosine_scheduler = CosineAnnealingLR(self.opt, T_max=main_steps)
+            self.scheduler = SequentialLR(self.opt, schedulers=[warm_up_scheduler, cosine_scheduler], milestones=[warm_up_steps])
+
+            logger = logging.getLogger()
+            time.sleep(0.5)
+            stage_time = datetime.now()
+            logger.info(f'Beta Model Training Start: {stage_time} seconds')
+
+
+            for epoch in range(num_epochs):
+                beta_loss_all = []
+                batch_shuffled_idx = np.random.permutation(self.data.shape[0])
+
+                time.sleep(0.5)
+                stage_time_last = datetime.now()
+                logger.info(f'Epoch {epoch + 1} Start: {stage_time_last} seconds')
+
+
+                for i in tqdm(range(interval)):
+
+                    start_pt = i * batch_size
+                    end_pt = min((i + 1) * batch_size, self.data.shape[0])
+                    local_idx = batch_shuffled_idx[start_pt:end_pt]
+                    batch = self.data[local_idx, ...]
+                    batch_votes = self.votes_data[local_idx, ...]
+
+                    comp = torch.sigmoid(batch_votes - self.votes_data.T)
+                    pred_comp = self.model(batch)
+
+                    beta_loss = torch.mean((comp - pred_comp) ** 2)
+
+                    beta_loss_all.append(beta_loss)
+
+                    self.opt.zero_grad()
+                    beta_loss.backward()
+                    self.opt.step()
+                    self.scheduler.step()  # Update LR after each optimizer step
+
+                beta_loss_all = torch.stack(beta_loss_all, dim=0)
+                print("iteration:", epoch + 1)
+                print("mean_beta_loss_all:", torch.mean(beta_loss_all).item())
+
+                time.sleep(0.5)
+                stage_time_last = datetime.now()
+                logger.info(f'Epoch {epoch + 1} End: {stage_time_last} seconds')
+
+                if save_dir is not None and (((epoch + 1) % 50 == 0) or ((epoch + 1) == num_epochs)):
+                    tmp_save_dir = Path(save_dir) / f'itr_{epoch + 1}'
+                    tmp_save_dir.mkdir(parents=True, exist_ok=True)
+                    model_file = tmp_save_dir / 'beta_model.pth'
+                    self.save_model(model_file)
+            self.ref_model = copy.deepcopy(self.model)
+        else:
+            self.load_model(load_dir)
+            self.ref_model = copy.deepcopy(self.model)
+
+    def online_update(self, dataset, num_epochs=1, warm_up_epochs=0, batch_size=1, lr = 1.0e-6):
+        interval = math.ceil(dataset["obs"].shape[0] / batch_size)
+        total_steps = num_epochs * interval
+        warm_up_steps = warm_up_epochs * interval
+        main_steps = total_steps - warm_up_steps
+
+        # Learning rate schedulers
+        self.lr = lr
+        self.opt = torch.optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=1e-5)
+        warm_up_scheduler = LinearLR(self.opt, start_factor=1e-8, end_factor=1.0, total_iters=warm_up_steps)
+        cosine_scheduler = CosineAnnealingLR(self.opt, T_max=main_steps)
+        self.scheduler = SequentialLR(self.opt, schedulers=[warm_up_scheduler, cosine_scheduler], milestones=[warm_up_steps])
+
+        logger = logging.getLogger()
+        time.sleep(0.5)
+        stage_time = datetime.now()
+        logger.info(f'Beta Model Training Start: {stage_time} seconds')
+
+        for epoch in range(num_epochs):
+            beta_loss_all = []
+            batch_shuffled_idx = np.random.permutation(dataset["obs"].shape[0])
+
+            time.sleep(0.5)
+            stage_time_last = datetime.now()
+            logger.info(f'Epoch {epoch + 1} Start: {stage_time_last} seconds')
+
+            for i in tqdm(range(interval)):
+
+                start_pt = i * batch_size
+                end_pt = min((i + 1) * batch_size, dataset["obs"].shape[0])
+                batch = index_batch(dataset, batch_shuffled_idx[start_pt:end_pt])
+
+                obs_1 = batch['obs']  # batch_size * traj_len * obs_dim
+                act_1 = batch['action']  # batch_size * traj_len * action_dim
+                obs_2 = batch['obs_2']
+                act_2 = batch['action_2']
+                s_a_1 = np.concatenate([obs_1, act_1], axis=-1)
+                s_a_2 = np.concatenate([obs_2, act_2], axis=-1)
+
+                votes_1 = torch.from_numpy(batch['votes']).to(self.device)
+                votes_2 = torch.from_numpy(batch['votes_2']).to(self.device)
+
+                # threshold = 1e-3
+                # diff = torch.abs(votes_1 - votes_2)
+                # condition_1 = (votes_1 > votes_2) & (diff >= threshold)  # votes_1 > votes_2 and diff >= threshold
+                # condition_2 = (votes_1 < votes_2) & (diff >= threshold)  # votes_1 < votes_2 and diff >= threshold
+
+                # comp_1 = torch.where(condition_1, torch.tensor(1.0, device=self.device), torch.tensor(0.0, device=self.device))
+                # comp_1 = torch.squeeze(comp_1, dim=-1)
+                # comp_2 = torch.where(condition_2, torch.tensor(1.0, device=self.device), torch.tensor(0.0, device=self.device))
+                # comp_2 = torch.squeeze(comp_2, dim=-1)
+
+                comp_1 = torch.sigmoid(votes_1 - votes_2)
+                comp_2 = torch.sigmoid(votes_2 - votes_1)
+
+                pred_comp_1 = self.model.one_to_one_forward(torch.from_numpy(s_a_1).float().to(self.device), torch.from_numpy(s_a_2).float().to(self.device))
+                pred_comp_2 = self.model.one_to_one_forward(torch.from_numpy(s_a_2).float().to(self.device), torch.from_numpy(s_a_1).float().to(self.device))
+                # ref_comp_1 = self.ref_model.one_to_one_forward(torch.from_numpy(s_a_1).float().to(self.device), torch.from_numpy(s_a_2).float().to(self.device)).detach()
+                # ref_comp_2 = self.ref_model.one_to_one_forward(torch.from_numpy(s_a_2).float().to(self.device), torch.from_numpy(s_a_1).float().to(self.device)).detach()
+
+                beta_loss = (torch.mean((comp_1 - pred_comp_1) ** 2) + torch.mean((comp_2 - pred_comp_2) ** 2)) / 2
+
+                beta_loss_all.append(beta_loss)
+
+                self.opt.zero_grad()
+                beta_loss.backward()
+                self.opt.step()
+                self.scheduler.step()  # Update LR after each optimizer step
+
+
+            time.sleep(0.5)
+            stage_time_now = datetime.now()
+            logger.info(f'Epoch {epoch + 1} Spending time:{stage_time_now - stage_time_last} End: {stage_time_now} seconds')
+
+            beta_loss_all = torch.stack(beta_loss_all, dim=0)
+            print("iteration:", epoch + 1)
+            print("mean_beta_loss_all:", torch.mean(beta_loss_all).item())
+
+    def save_model(self, filepath):
+        torch.save(self.state_dict(), filepath)
+
+    def load_model(self, filepath):
+        self.load_state_dict(torch.load(filepath, map_location=self.device))
+
+
+class BetaImageNetwork(nn.Module):
+    def __init__(self, data, policy: DiffusionTransformerHybridImagePolicy, device=torch.device('cuda'), data_size = 500, ):
+        super(BetaNetwork, self).__init__()
+
+        act_data = np.concatenate((data['action'], data['action_2']), axis=0)
+        votes_data = np.concatenate((data['votes'], data['votes_2']), axis=0)
+        obs_data = {key:np.concatenate((data['obs'][key], data['obs_2'][key]), axis=0) for key in data['obs']}
+        if data_size <= act_data.shape[0]:
+            indices = np.random.randint(0, act_data.shape[0], size=data_size)
+            obs_data = {key:obs_data[key][indices, ...]for key in obs_data.keys()}
+            act_data = act_data[indices, ...]
+            votes_data = votes_data[indices, ...]
+
+        act_data = torch.from_numpy(act_data).float().to(device)
+        obs_data = {key: torch.from_numpy(obs_data[key]).float().to(device) for key in obs_data.keys()}
+
+        obs_encoder = policy.nets['policy'].nets['encoder'].nets['obs']
+
+        obs_data = policy.normalizer.normalize(obs_data)
+        act_data = policy.normalizer.normalize(act_data)
+
+        this_nobs = dict_apply(obs_data, 
+                    lambda x: x.reshape(-1,*x.shape[2:]))
+        nobs_features = obs_encoder(this_nobs)
+        nobs_features = nobs_features.reshape(data_size, act_data.shape[1], -1)
+        nobs_features.detach()
+
+
+        self.votes_data = torch.from_numpy(votes_data).to(device)
+        self.lr = None
+        self.device = device
+        self.data = torch.concat((act_data, nobs_features), dim=-1)
+
+        class BetaModel(nn.Module):
+            def __init__(self, obs_data, act_data, device=torch.device('cuda')):
+                super(BetaModel, self).__init__()
+
+                self.enc_model = TransformerEncModel(
+                    data_dim = act_data.shape[-1] + obs_data.shape[-1],
+                    embedding_dim = 256,
+                    nhead = 4,
+                    num_encoder_layers = 2,
+                    device = device
+                ).to(device)
+
+                self.comp_model = AttentionComparisonModel(
+                    input_dim = 256, 
+                    dropout_rate = 0.3,
+                    nhead = 16,
+                    device = device
+                ).to(device)
+
+                self.data = torch.concat((obs_data, act_data), dim=-1)
+
+            def forward(self, x):
+                batch_f = self.enc_model(x)
+                all_data_f = self.enc_model(self.data)
+                bias = all_data_f.mean()
+                std = all_data_f.std()
+                all_data_f = (all_data_f - bias) / std
+                batch_f = (batch_f - bias) / std
+                output = self.comp_model(batch_f, all_data_f)
+                return output
+            
+            def one_to_one_forward(self, x, y):
+                x_f = self.enc_model(x)
+                y_f = self.enc_model(y)
+                all_data_f = self.enc_model(self.data)
+                bias = all_data_f.mean()
+                std = all_data_f.std()
+                x_f = (x_f - bias) / std
+                y_f = (y_f - bias) / std
+                output = self.comp_model.one_to_one_forward(x_f, y_f)
+                return output
+
+       
+
+        self.model = BetaModel(nobs_features, act_data, device)
+        self.opt = None #, weight_decay=1.0e-4
+        self.scheduler = None #torch.optim.lr_scheduler.ReduceLROnPlateau(self.opt, mode='min', patience=100, verbose=True)
+        self.ref_model: BetaModel
+
+    def get_alpha_beta(self, x):
+        batch_comp = self.model(x).detach()
+        alpha = torch.sum(batch_comp, dim=-1)
+        beta = torch.sum(1 - batch_comp, dim=-1)
+        # alpha = torch.sum(torch.where(batch_comp > 0.5, torch.tensor(1.0, device=self.device), torch.tensor(0.0, device=self.device)), dim=-1)
+        # beta = torch.sum(torch.where(batch_comp < 0.5, torch.tensor(1.0, device=self.device), torch.tensor(0.0, device=self.device)), dim=-1)
+
+        return alpha.detach(), beta.detach()
+
+    # def fit_data(self, save_dir=None, load_dir=None, num_epochs=1, warm_up_epochs=0, batch_size=1, lr=1.0e-5):
+    #     if load_dir is None:
+    #         interval = math.ceil(self.data.shape[0] / batch_size)
+    #         total_steps = num_epochs * interval
+    #         warm_up_steps = warm_up_epochs * interval
+    #         main_steps = total_steps - warm_up_steps
+
+    #         # Learning rate schedulers
+    #         self.lr = lr
+    #         self.opt = torch.optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=1e-5)
+    #         warm_up_scheduler = LinearLR(self.opt, start_factor=1e-8, end_factor=1.0, total_iters=warm_up_steps)
+    #         cosine_scheduler = CosineAnnealingLR(self.opt, T_max=main_steps)
+    #         self.scheduler = SequentialLR(self.opt, schedulers=[warm_up_scheduler, cosine_scheduler], milestones=[warm_up_steps])
+
+    #         for epoch in range(num_epochs):
+    #             beta_loss_all = []
+
+    #             batch_shuffled_idx = np.random.permutation(self.data.shape[0])
+    #             for i in tqdm(range(interval)):
+
+    #                 start_pt = i * batch_size
+    #                 end_pt = min((i + 1) * batch_size, self.data.shape[0])
+    #                 local_idx = batch_shuffled_idx[start_pt:end_pt]
+    #                 batch = self.data[local_idx, ...]
+    #                 batch_votes = self.votes_data[local_idx, ...]
+
+    #                 comp = torch.sigmoid(batch_votes - self.votes_data.T)
+    #                 pred_comp = self.model(batch)
+
+    #                 beta_loss = torch.mean((comp - pred_comp) ** 2)
+
+    #                 beta_loss_all.append(beta_loss)
+
+    #                 self.opt.zero_grad()
+    #                 beta_loss.backward()
+    #                 self.opt.step()
+    #                 self.scheduler.step()  # Update LR after each optimizer step
+
+    #             beta_loss_all = torch.stack(beta_loss_all, dim=0)
+    #             print("iteration:", epoch + 1)
+    #             print("mean_beta_loss_all:", torch.mean(beta_loss_all).item())
+
+    #             if save_dir is not None and (((epoch + 1) % 50 == 0) or ((epoch + 1) == num_epochs)):
+    #                 tmp_save_dir = Path(save_dir) / f'itr_{epoch + 1}'
+    #                 tmp_save_dir.mkdir(parents=True, exist_ok=True)
+    #                 model_file = tmp_save_dir / 'beta_model.pth'
+    #                 self.save_model(model_file)
+    #         self.ref_model = copy.deepcopy(self.model)
+    #     else:
+    #         self.load_model(load_dir)
+    #         self.ref_model = copy.deepcopy(self.model)
+
+    def online_update(self, dataset, num_epochs=1, warm_up_epochs=0, batch_size=1, lr = 1.0e-6):
+        interval = math.ceil(dataset["obs"].shape[0] / batch_size)
+        total_steps = num_epochs * interval
+        warm_up_steps = warm_up_epochs * interval
+        main_steps = total_steps - warm_up_steps
+
+        # Learning rate schedulers
+        self.lr = lr
+        self.opt = torch.optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=1e-5)
+        warm_up_scheduler = LinearLR(self.opt, start_factor=1e-8, end_factor=1.0, total_iters=warm_up_steps)
+        cosine_scheduler = CosineAnnealingLR(self.opt, T_max=main_steps)
+        self.scheduler = SequentialLR(self.opt, schedulers=[warm_up_scheduler, cosine_scheduler], milestones=[warm_up_steps])
+
+        for epoch in range(num_epochs):
+            beta_loss_all = []
+
+            batch_shuffled_idx = np.random.permutation(dataset["obs"].shape[0])
+            for i in tqdm(range(interval)):
+
+                start_pt = i * batch_size
+                end_pt = min((i + 1) * batch_size, dataset["obs"].shape[0])
+                batch = index_batch(dataset, batch_shuffled_idx[start_pt:end_pt])
+
+                obs_1 = batch['obs']  # batch_size * traj_len * obs_dim
+                act_1 = batch['action']  # batch_size * traj_len * action_dim
+                obs_2 = batch['obs_2']
+                act_2 = batch['action_2']
+                s_a_1 = np.concatenate([obs_1, act_1], axis=-1)
+                s_a_2 = np.concatenate([obs_2, act_2], axis=-1)
+
+                votes_1 = torch.from_numpy(batch['votes']).to(self.device)
+                votes_2 = torch.from_numpy(batch['votes_2']).to(self.device)
+
+                # threshold = 1e-3
+                # diff = torch.abs(votes_1 - votes_2)
+                # condition_1 = (votes_1 > votes_2) & (diff >= threshold)  # votes_1 > votes_2 and diff >= threshold
+                # condition_2 = (votes_1 < votes_2) & (diff >= threshold)  # votes_1 < votes_2 and diff >= threshold
+
+                # comp_1 = torch.where(condition_1, torch.tensor(1.0, device=self.device), torch.tensor(0.0, device=self.device))
+                # comp_1 = torch.squeeze(comp_1, dim=-1)
+                # comp_2 = torch.where(condition_2, torch.tensor(1.0, device=self.device), torch.tensor(0.0, device=self.device))
+                # comp_2 = torch.squeeze(comp_2, dim=-1)
+
+                comp_1 = torch.sigmoid(votes_1 - votes_2)
+                comp_2 = torch.sigmoid(votes_2 - votes_1)
+
+                pred_comp_1 = self.model.one_to_one_forward(torch.from_numpy(s_a_1).float().to(self.device), torch.from_numpy(s_a_2).float().to(self.device))
+                pred_comp_2 = self.model.one_to_one_forward(torch.from_numpy(s_a_2).float().to(self.device), torch.from_numpy(s_a_1).float().to(self.device))
+                # ref_comp_1 = self.ref_model.one_to_one_forward(torch.from_numpy(s_a_1).float().to(self.device), torch.from_numpy(s_a_2).float().to(self.device)).detach()
+                # ref_comp_2 = self.ref_model.one_to_one_forward(torch.from_numpy(s_a_2).float().to(self.device), torch.from_numpy(s_a_1).float().to(self.device)).detach()
+
+                beta_loss = (torch.mean((comp_1 - pred_comp_1) ** 2) + torch.mean((comp_2 - pred_comp_2) ** 2)) / 2
+
+                beta_loss_all.append(beta_loss)
+
+                self.opt.zero_grad()
+                beta_loss.backward()
+                self.opt.step()
+                self.scheduler.step()  # Update LR after each optimizer step
+
+            beta_loss_all = torch.stack(beta_loss_all, dim=0)
+            print("iteration:", epoch + 1)
+            print("mean_beta_loss_all:", torch.mean(beta_loss_all).item())
+
+    def save_model(self, filepath):
+        torch.save(self.state_dict(), filepath)
+
+    def load_model(self, filepath):
+        self.load_state_dict(torch.load(filepath, map_location=self.device))

common/pymunk_override.py

@@ -0,0 +1,248 @@
+# ----------------------------------------------------------------------------
+# pymunk
+# Copyright (c) 2007-2016 Victor Blomqvist
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# ----------------------------------------------------------------------------
+
+"""This submodule contains helper functions to help with quick prototyping 
+using pymunk together with pygame.
+
+Intended to help with debugging and prototyping, not for actual production use
+in a full application. The methods contained in this module is opinionated 
+about your coordinate system and not in any way optimized. 
+"""
+
+__docformat__ = "reStructuredText"
+
+__all__ = [
+    "DrawOptions",
+    "get_mouse_pos",
+    "to_pygame",
+    "from_pygame",
+    "lighten",
+    "positive_y_is_up",
+]
+
+from typing import List, Sequence, Tuple
+
+import pygame
+
+import numpy as np
+
+import pymunk
+from pymunk.space_debug_draw_options import SpaceDebugColor
+from pymunk.vec2d import Vec2d
+
+positive_y_is_up: bool = False
+"""Make increasing values of y point upwards.
+
+When True::
+
+    y
+    ^
+    |      . (3, 3)
+    |
+    |   . (2, 2)
+    |
+    +------ > x
+    
+When False::
+
+    +------ > x
+    |
+    |   . (2, 2)
+    |
+    |      . (3, 3)
+    v
+    y
+    
+"""
+
+
+class DrawOptions(pymunk.SpaceDebugDrawOptions):
+    def __init__(self, surface: pygame.Surface) -> None:
+        """Draw a pymunk.Space on a pygame.Surface object.
+
+        Typical usage::
+
+        >>> import pymunk
+        >>> surface = pygame.Surface((10,10))
+        >>> space = pymunk.Space()
+        >>> options = pymunk.pygame_util.DrawOptions(surface)
+        >>> space.debug_draw(options)
+
+        You can control the color of a shape by setting shape.color to the color
+        you want it drawn in::
+
+        >>> c = pymunk.Circle(None, 10)
+        >>> c.color = pygame.Color("pink")
+
+        See pygame_util.demo.py for a full example
+
+        Since pygame uses a coordinate system where y points down (in contrast
+        to many other cases), you either have to make the physics simulation
+        with Pymunk also behave in that way, or flip everything when you draw.
+
+        The easiest is probably to just make the simulation behave the same
+        way as Pygame does. In that way all coordinates used are in the same
+        orientation and easy to reason about::
+
+        >>> space = pymunk.Space()
+        >>> space.gravity = (0, -1000)
+        >>> body = pymunk.Body()
+        >>> body.position = (0, 0) # will be positioned in the top left corner
+        >>> space.debug_draw(options)
+
+        To flip the drawing its possible to set the module property
+        :py:data:`positive_y_is_up` to True. Then the pygame drawing will flip
+        the simulation upside down before drawing::
+
+        >>> positive_y_is_up = True
+        >>> body = pymunk.Body()
+        >>> body.position = (0, 0)
+        >>> # Body will be position in bottom left corner
+
+        :Parameters:
+                surface : pygame.Surface
+                    Surface that the objects will be drawn on
+        """
+        self.surface = surface
+        super(DrawOptions, self).__init__()
+
+    def draw_circle(
+        self,
+        pos: Vec2d,
+        angle: float,
+        radius: float,
+        outline_color: SpaceDebugColor,
+        fill_color: SpaceDebugColor,
+    ) -> None:
+        p = to_pygame(pos, self.surface)
+
+        pygame.draw.circle(self.surface, fill_color.as_int(), p, round(radius), 0)
+        pygame.draw.circle(self.surface, light_color(fill_color).as_int(), p, round(radius-4), 0)
+
+        circle_edge = pos + Vec2d(radius, 0).rotated(angle)
+        p2 = to_pygame(circle_edge, self.surface)
+        line_r = 2 if radius > 20 else 1
+        # pygame.draw.lines(self.surface, outline_color.as_int(), False, [p, p2], line_r)
+
+    def draw_segment(self, a: Vec2d, b: Vec2d, color: SpaceDebugColor) -> None:
+        p1 = to_pygame(a, self.surface)
+        p2 = to_pygame(b, self.surface)
+
+        pygame.draw.aalines(self.surface, color.as_int(), False, [p1, p2])
+
+    def draw_fat_segment(
+        self,
+        a: Tuple[float, float],
+        b: Tuple[float, float],
+        radius: float,
+        outline_color: SpaceDebugColor,
+        fill_color: SpaceDebugColor,
+    ) -> None:
+        p1 = to_pygame(a, self.surface)
+        p2 = to_pygame(b, self.surface)
+
+        r = round(max(1, radius * 2))
+        pygame.draw.lines(self.surface, fill_color.as_int(), False, [p1, p2], r)
+        if r > 2:
+            orthog = [abs(p2[1] - p1[1]), abs(p2[0] - p1[0])]
+            if orthog[0] == 0 and orthog[1] == 0:
+                return
+            scale = radius / (orthog[0] * orthog[0] + orthog[1] * orthog[1]) ** 0.5
+            orthog[0] = round(orthog[0] * scale)
+            orthog[1] = round(orthog[1] * scale)
+            points = [
+                (p1[0] - orthog[0], p1[1] - orthog[1]),
+                (p1[0] + orthog[0], p1[1] + orthog[1]),
+                (p2[0] + orthog[0], p2[1] + orthog[1]),
+                (p2[0] - orthog[0], p2[1] - orthog[1]),
+            ]
+            pygame.draw.polygon(self.surface, fill_color.as_int(), points)
+            pygame.draw.circle(
+                self.surface,
+                fill_color.as_int(),
+                (round(p1[0]), round(p1[1])),
+                round(radius),
+            )
+            pygame.draw.circle(
+                self.surface,
+                fill_color.as_int(),
+                (round(p2[0]), round(p2[1])),
+                round(radius),
+            )
+
+    def draw_polygon(
+        self,
+        verts: Sequence[Tuple[float, float]],
+        radius: float,
+        outline_color: SpaceDebugColor,
+        fill_color: SpaceDebugColor,
+    ) -> None:
+        ps = [to_pygame(v, self.surface) for v in verts]
+        ps += [ps[0]]
+
+        radius = 2
+        pygame.draw.polygon(self.surface, light_color(fill_color).as_int(), ps)
+
+        if radius > 0:
+            for i in range(len(verts)):
+                a = verts[i]
+                b = verts[(i + 1) % len(verts)]
+                self.draw_fat_segment(a, b, radius, fill_color, fill_color)
+
+    def draw_dot(
+        self, size: float, pos: Tuple[float, float], color: SpaceDebugColor
+    ) -> None:
+        p = to_pygame(pos, self.surface)
+        pygame.draw.circle(self.surface, color.as_int(), p, round(size), 0)
+
+
+def get_mouse_pos(surface: pygame.Surface) -> Tuple[int, int]:
+    """Get position of the mouse pointer in pymunk coordinates."""
+    p = pygame.mouse.get_pos()
+    return from_pygame(p, surface)
+
+
+def to_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]:
+    """Convenience method to convert pymunk coordinates to pygame surface
+    local coordinates.
+
+    Note that in case positive_y_is_up is False, this function won't actually do
+    anything except converting the point to integers.
+    """
+    if positive_y_is_up:
+        return round(p[0]), surface.get_height() - round(p[1])
+    else:
+        return round(p[0]), round(p[1])
+
+
+def from_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]:
+    """Convenience method to convert pygame surface local coordinates to
+    pymunk coordinates
+    """
+    return to_pygame(p, surface)
+
+
+def light_color(color: SpaceDebugColor):
+    color = np.minimum(1.2 * np.float32([color.r, color.g, color.b, color.a]), np.float32([255]))
+    color = SpaceDebugColor(r=color[0], g=color[1], b=color[2], a=color[3])
+    return color

common/pymunk_util.py

@@ -0,0 +1,52 @@
+import pygame
+import pymunk
+import pymunk.pygame_util
+import numpy as np
+
+COLLTYPE_DEFAULT = 0
+COLLTYPE_MOUSE = 1
+COLLTYPE_BALL = 2
+
+def get_body_type(static=False):
+    body_type = pymunk.Body.DYNAMIC
+    if static:
+        body_type = pymunk.Body.STATIC
+    return body_type
+
+
+def create_rectangle(space,
+        pos_x,pos_y,width,height,
+        density=3,static=False):
+    body = pymunk.Body(body_type=get_body_type(static))
+    body.position = (pos_x,pos_y)
+    shape = pymunk.Poly.create_box(body,(width,height))
+    shape.density = density
+    space.add(body,shape)
+    return body, shape
+
+
+def create_rectangle_bb(space, 
+        left, bottom, right, top, 
+        **kwargs):
+    pos_x = (left + right) / 2
+    pos_y = (top + bottom) / 2
+    height = top - bottom
+    width = right - left
+    return create_rectangle(space, pos_x, pos_y, width, height, **kwargs)
+
+def create_circle(space, pos_x, pos_y, radius, density=3, static=False):
+    body = pymunk.Body(body_type=get_body_type(static))
+    body.position = (pos_x, pos_y)
+    shape = pymunk.Circle(body, radius=radius)
+    shape.density = density
+    shape.collision_type = COLLTYPE_BALL
+    space.add(body, shape)
+    return body, shape
+
+def get_body_state(body):
+    state = np.zeros(6, dtype=np.float32)
+    state[:2] = body.position
+    state[2] = body.angle
+    state[3:5] = body.velocity
+    state[5] = body.angular_velocity
+    return state

common/pytorch_util.py

@@ -0,0 +1,82 @@
+from typing import Dict, Callable, List
+import collections
+import torch
+import torch.nn as nn
+
+def dict_apply(
+        x: Dict[str, torch.Tensor], 
+        func: Callable[[torch.Tensor], torch.Tensor]
+        ) -> Dict[str, torch.Tensor]:
+    result = dict()
+    for key, value in x.items():
+        if isinstance(value, dict):
+            result[key] = dict_apply(value, func)
+        else:
+            result[key] = func(value)
+    return result
+
+def pad_remaining_dims(x, target):
+    assert x.shape == target.shape[:len(x.shape)]
+    return x.reshape(x.shape + (1,)*(len(target.shape) - len(x.shape)))
+
+def dict_apply_split(
+        x: Dict[str, torch.Tensor], 
+        split_func: Callable[[torch.Tensor], Dict[str, torch.Tensor]]
+        ) -> Dict[str, torch.Tensor]:
+    results = collections.defaultdict(dict)
+    for key, value in x.items():
+        result = split_func(value)
+        for k, v in result.items():
+            results[k][key] = v
+    return results
+
+def dict_apply_reduce(
+        x: List[Dict[str, torch.Tensor]],
+        reduce_func: Callable[[List[torch.Tensor]], torch.Tensor]
+        ) -> Dict[str, torch.Tensor]:
+    result = dict()
+    for key in x[0].keys():
+        result[key] = reduce_func([x_[key] for x_ in x])
+    return result
+
+
+def replace_submodules(
+        root_module: nn.Module, 
+        predicate: Callable[[nn.Module], bool], 
+        func: Callable[[nn.Module], nn.Module]) -> nn.Module:
+    """
+    predicate: Return true if the module is to be replaced.
+    func: Return new module to use.
+    """
+    if predicate(root_module):
+        return func(root_module)
+
+    bn_list = [k.split('.') for k, m 
+        in root_module.named_modules(remove_duplicate=True) 
+        if predicate(m)]
+    for *parent, k in bn_list:
+        parent_module = root_module
+        if len(parent) > 0:
+            parent_module = root_module.get_submodule('.'.join(parent))
+        if isinstance(parent_module, nn.Sequential):
+            src_module = parent_module[int(k)]
+        else:
+            src_module = getattr(parent_module, k)
+        tgt_module = func(src_module)
+        if isinstance(parent_module, nn.Sequential):
+            parent_module[int(k)] = tgt_module
+        else:
+            setattr(parent_module, k, tgt_module)
+    # verify that all BN are replaced
+    bn_list = [k.split('.') for k, m 
+        in root_module.named_modules(remove_duplicate=True) 
+        if predicate(m)]
+    assert len(bn_list) == 0
+    return root_module
+
+def optimizer_to(optimizer, device):
+    for state in optimizer.state.values():
+        for k, v in state.items():
+            if isinstance(v, torch.Tensor):
+                state[k] = v.to(device=device)
+    return optimizer

common/replay_buffer.py

@@ -0,0 +1,615 @@
+from typing import Union, Dict, Optional
+import os
+import math
+import numbers
+import zarr
+import numcodecs
+import numpy as np
+from functools import cached_property
+
+def check_chunks_compatible(chunks: tuple, shape: tuple):
+    assert len(shape) == len(chunks)
+    for c in chunks:
+        assert isinstance(c, numbers.Integral)
+        assert c > 0
+
+def rechunk_recompress_array(group, name, 
+        chunks=None, chunk_length=None,
+        compressor=None, tmp_key='_temp'):
+    old_arr = group[name]
+    if chunks is None:
+        if chunk_length is not None:
+            chunks = (chunk_length,) + old_arr.chunks[1:]
+        else:
+            chunks = old_arr.chunks
+    check_chunks_compatible(chunks, old_arr.shape)
+    
+    if compressor is None:
+        compressor = old_arr.compressor
+    
+    if (chunks == old_arr.chunks) and (compressor == old_arr.compressor):
+        # no change
+        return old_arr
+
+    # rechunk recompress
+    group.move(name, tmp_key)
+    old_arr = group[tmp_key]
+    n_copied, n_skipped, n_bytes_copied = zarr.copy(
+        source=old_arr,
+        dest=group,
+        name=name,
+        chunks=chunks,
+        compressor=compressor,
+    )
+    del group[tmp_key]
+    arr = group[name]
+    return arr
+
+def get_optimal_chunks(shape, dtype, 
+        target_chunk_bytes=2e6, 
+        max_chunk_length=None):
+    """
+    Common shapes
+    T,D
+    T,N,D
+    T,H,W,C
+    T,N,H,W,C
+    """
+    itemsize = np.dtype(dtype).itemsize
+    # reversed
+    rshape = list(shape[::-1])
+    if max_chunk_length is not None:
+        rshape[-1] = int(max_chunk_length)
+    split_idx = len(shape)-1
+    for i in range(len(shape)-1):
+        this_chunk_bytes = itemsize * np.prod(rshape[:i])
+        next_chunk_bytes = itemsize * np.prod(rshape[:i+1])
+        if this_chunk_bytes <= target_chunk_bytes \
+            and next_chunk_bytes > target_chunk_bytes:
+            split_idx = i
+
+    rchunks = rshape[:split_idx]
+    item_chunk_bytes = itemsize * np.prod(rshape[:split_idx])
+    this_max_chunk_length = rshape[split_idx]
+    next_chunk_length = min(this_max_chunk_length, math.ceil(
+            target_chunk_bytes / item_chunk_bytes))
+    rchunks.append(next_chunk_length)
+    len_diff = len(shape) - len(rchunks)
+    rchunks.extend([1] * len_diff)
+    chunks = tuple(rchunks[::-1])
+    # print(np.prod(chunks) * itemsize / target_chunk_bytes)
+    return chunks
+
+
+class ReplayBuffer:
+    """
+    Zarr-based temporal datastructure.
+    Assumes first dimension to be time. Only chunk in time dimension.
+    """
+    def __init__(self, 
+            root: Union[zarr.Group, 
+            Dict[str,dict]]):
+        """
+        Dummy constructor. Use copy_from* and create_from* class methods instead.
+        """
+        assert('data' in root)
+        assert('meta' in root)
+        assert('episode_ends' in root['meta'])
+        for key, value in root['data'].items():
+            assert(value.shape[0] == root['meta']['episode_ends'][-1])
+        self.root = root
+    
+    # ============= create constructors ===============
+    @classmethod
+    def create_empty_zarr(cls, storage=None, root=None):
+        if root is None:
+            if storage is None:
+                storage = zarr.MemoryStore()
+            root = zarr.group(store=storage)
+        data = root.require_group('data', overwrite=False)
+        meta = root.require_group('meta', overwrite=False)
+        if 'episode_ends' not in meta:
+            episode_ends = meta.zeros('episode_ends', shape=(0,), dtype=np.int64,
+                compressor=None, overwrite=False)
+        return cls(root=root)
+    
+    @classmethod
+    def create_empty_numpy(cls):
+        root = {
+            'data': dict(),
+            'meta': {
+                'episode_ends': np.zeros((0,), dtype=np.int64)
+            }
+        }
+        return cls(root=root)
+    
+    @classmethod
+    def create_from_group(cls, group, **kwargs):
+        if 'data' not in group:
+            # create from stratch
+            buffer = cls.create_empty_zarr(root=group, **kwargs)
+        else:
+            # already exist
+            buffer = cls(root=group, **kwargs)
+        return buffer
+
+    @classmethod
+    def create_from_path(cls, zarr_path, mode='r', **kwargs):
+        """
+        Open a on-disk zarr directly (for dataset larger than memory).
+        Slower.
+        """
+        group = zarr.open(os.path.expanduser(zarr_path), mode)
+        return cls.create_from_group(group, **kwargs)
+    
+    # ============= copy constructors ===============
+    @classmethod
+    def copy_from_store(cls, src_store, store=None, keys=None, 
+            chunks: Dict[str,tuple]=dict(), 
+            compressors: Union[dict, str, numcodecs.abc.Codec]=dict(), 
+            if_exists='replace',
+            **kwargs):
+        """
+        Load to memory.
+        """
+        src_root = zarr.group(src_store)
+        root = None
+        if store is None:
+            # numpy backend
+            meta = dict()
+            for key, value in src_root['meta'].items():
+                if len(value.shape) == 0:
+                    meta[key] = np.array(value)
+                else:
+                    meta[key] = value[:]
+
+            if keys is None:
+                keys = src_root['data'].keys()
+            data = dict()
+            for key in keys:
+                arr = src_root['data'][key]
+                data[key] = arr[:]
+
+            root = {
+                'meta': meta,
+                'data': data
+            }
+        else:
+            root = zarr.group(store=store)
+            # copy without recompression
+            n_copied, n_skipped, n_bytes_copied = zarr.copy_store(source=src_store, dest=store,
+                source_path='/meta', dest_path='/meta', if_exists=if_exists)
+            data_group = root.create_group('data', overwrite=True)
+            if keys is None:
+                keys = src_root['data'].keys()
+            for key in keys:
+                value = src_root['data'][key]
+                cks = cls._resolve_array_chunks(
+                    chunks=chunks, key=key, array=value)
+                cpr = cls._resolve_array_compressor(
+                    compressors=compressors, key=key, array=value)
+                if cks == value.chunks and cpr == value.compressor:
+                    # copy without recompression
+                    this_path = '/data/' + key
+                    n_copied, n_skipped, n_bytes_copied = zarr.copy_store(
+                        source=src_store, dest=store,
+                        source_path=this_path, dest_path=this_path,
+                        if_exists=if_exists
+                    )
+                else:
+                    # copy with recompression
+                    n_copied, n_skipped, n_bytes_copied = zarr.copy(
+                        source=value, dest=data_group, name=key,
+                        chunks=cks, compressor=cpr, if_exists=if_exists
+                    )
+        buffer = cls(root=root)
+        return buffer
+    
+    @classmethod
+    def copy_from_path(cls, zarr_path, backend=None, store=None, keys=None, 
+            chunks: Dict[str,tuple]=dict(), 
+            compressors: Union[dict, str, numcodecs.abc.Codec]=dict(), 
+            if_exists='replace',
+            **kwargs):
+        """
+        Copy a on-disk zarr to in-memory compressed.
+        Recommended
+        """
+        if backend == 'numpy':
+            print('backend argument is deprecated!')
+            store = None
+        group = zarr.open(os.path.expanduser(zarr_path), 'r')
+        return cls.copy_from_store(src_store=group.store, store=store, 
+            keys=keys, chunks=chunks, compressors=compressors, 
+            if_exists=if_exists, **kwargs)
+
+    # ============= save methods ===============
+    def save_to_store(self, store, 
+            chunks: Optional[Dict[str,tuple]]=dict(),
+            compressors: Union[str, numcodecs.abc.Codec, dict]=dict(),
+            if_exists='replace', 
+            **kwargs):
+        
+        root = zarr.group(store)
+        if self.backend == 'zarr':
+            # recompression free copy
+            n_copied, n_skipped, n_bytes_copied = zarr.copy_store(
+                source=self.root.store, dest=store,
+                source_path='/meta', dest_path='/meta', if_exists=if_exists)
+        else:
+            meta_group = root.create_group('meta', overwrite=True)
+            # save meta, no chunking
+            for key, value in self.root['meta'].items():
+                _ = meta_group.array(
+                    name=key,
+                    data=value, 
+                    shape=value.shape, 
+                    chunks=value.shape)
+        
+        # save data, chunk
+        data_group = root.create_group('data', overwrite=True)
+        for key, value in self.root['data'].items():
+            cks = self._resolve_array_chunks(
+                chunks=chunks, key=key, array=value)
+            cpr = self._resolve_array_compressor(
+                compressors=compressors, key=key, array=value)
+            if isinstance(value, zarr.Array):
+                if cks == value.chunks and cpr == value.compressor:
+                    # copy without recompression
+                    this_path = '/data/' + key
+                    n_copied, n_skipped, n_bytes_copied = zarr.copy_store(
+                        source=self.root.store, dest=store,
+                        source_path=this_path, dest_path=this_path, if_exists=if_exists)
+                else:
+                    # copy with recompression
+                    n_copied, n_skipped, n_bytes_copied = zarr.copy(
+                        source=value, dest=data_group, name=key,
+                        chunks=cks, compressor=cpr, if_exists=if_exists
+                    )
+            else:
+                # numpy
+                _ = data_group.array(
+                    name=key,
+                    data=value,
+                    chunks=cks,
+                    compressor=cpr
+                )
+        return store
+
+    def save_to_path(self, zarr_path,             
+            chunks: Optional[Dict[str,tuple]]=dict(),
+            compressors: Union[str, numcodecs.abc.Codec, dict]=dict(), 
+            if_exists='replace', 
+            **kwargs):
+        store = zarr.DirectoryStore(os.path.expanduser(zarr_path))
+        return self.save_to_store(store, chunks=chunks, 
+            compressors=compressors, if_exists=if_exists, **kwargs)
+
+    @staticmethod
+    def resolve_compressor(compressor='default'):
+        if compressor == 'default':
+            compressor = numcodecs.Blosc(cname='lz4', clevel=5, 
+                shuffle=numcodecs.Blosc.NOSHUFFLE)
+        elif compressor == 'disk':
+            compressor = numcodecs.Blosc('zstd', clevel=5, 
+                shuffle=numcodecs.Blosc.BITSHUFFLE)
+        return compressor
+
+    @classmethod
+    def _resolve_array_compressor(cls, 
+            compressors: Union[dict, str, numcodecs.abc.Codec], key, array):
+        # allows compressor to be explicitly set to None
+        cpr = 'nil'
+        if isinstance(compressors, dict):
+            if key in compressors:
+                cpr = cls.resolve_compressor(compressors[key])
+            elif isinstance(array, zarr.Array):
+                cpr = array.compressor
+        else:
+            cpr = cls.resolve_compressor(compressors)
+        # backup default
+        if cpr == 'nil':
+            cpr = cls.resolve_compressor('default')
+        return cpr
+    
+    @classmethod
+    def _resolve_array_chunks(cls,
+            chunks: Union[dict, tuple], key, array):
+        cks = None
+        if isinstance(chunks, dict):
+            if key in chunks:
+                cks = chunks[key]
+            elif isinstance(array, zarr.Array):
+                cks = array.chunks
+        elif isinstance(chunks, tuple):
+            cks = chunks
+        else:
+            raise TypeError(f"Unsupported chunks type {type(chunks)}")
+        # backup default
+        if cks is None:
+            cks = get_optimal_chunks(shape=array.shape, dtype=array.dtype)
+        # check
+        check_chunks_compatible(chunks=cks, shape=array.shape)
+        return cks
+    
+    # ============= properties =================
+    @cached_property
+    def data(self):
+        return self.root['data']
+    
+    @cached_property
+    def meta(self):
+        return self.root['meta']
+
+    def update_meta(self, data):
+        # sanitize data
+        np_data = dict()
+        for key, value in data.items():
+            if isinstance(value, np.ndarray):
+                np_data[key] = value
+            else:
+                arr = np.array(value)
+                if arr.dtype == object:
+                    raise TypeError(f"Invalid value type {type(value)}")
+                np_data[key] = arr
+
+        meta_group = self.meta
+        if self.backend == 'zarr':
+            for key, value in np_data.items():
+                _ = meta_group.array(
+                    name=key,
+                    data=value, 
+                    shape=value.shape, 
+                    chunks=value.shape,
+                    overwrite=True)
+        else:
+            meta_group.update(np_data)
+        
+        return meta_group
+    
+    @property
+    def episode_ends(self):
+        return self.meta['episode_ends']
+    
+    def get_episode_idxs(self):
+        import numba
+        numba.jit(nopython=True)
+        def _get_episode_idxs(episode_ends):
+            result = np.zeros((episode_ends[-1],), dtype=np.int64)
+            for i in range(len(episode_ends)):
+                start = 0
+                if i > 0:
+                    start = episode_ends[i-1]
+                end = episode_ends[i]
+                for idx in range(start, end):
+                    result[idx] = i
+            return result
+        return _get_episode_idxs(self.episode_ends)
+        
+    
+    @property
+    def backend(self):
+        backend = 'numpy'
+        if isinstance(self.root, zarr.Group):
+            backend = 'zarr'
+        return backend
+    
+    # =========== dict-like API ==============
+    def __repr__(self) -> str:
+        if self.backend == 'zarr':
+            return str(self.root.tree())
+        else:
+            return super().__repr__()
+
+    def keys(self):
+        return self.data.keys()
+    
+    def values(self):
+        return self.data.values()
+    
+    def items(self):
+        return self.data.items()
+    
+    def __getitem__(self, key):
+        return self.data[key]
+
+    def __contains__(self, key):
+        return key in self.data
+
+    # =========== our API ==============
+    @property
+    def n_steps(self):
+        if len(self.episode_ends) == 0:
+            return 0
+        return self.episode_ends[-1]
+    
+    @property
+    def n_episodes(self):
+        return len(self.episode_ends)
+
+    @property
+    def chunk_size(self):
+        if self.backend == 'zarr':
+            return next(iter(self.data.arrays()))[-1].chunks[0]
+        return None
+
+    @property
+    def episode_lengths(self):
+        ends = self.episode_ends[:]
+        ends = np.insert(ends, 0, 0)
+        lengths = np.diff(ends)
+        return lengths
+
+    def add_episode(self, 
+                    data: Dict[str, np.ndarray],
+                    meta_data: Optional[Dict[str, Union[np.ndarray, int]]] = None,
+                    chunks: Optional[Dict[str, tuple]] = dict(),
+                    compressors: Union[str, numcodecs.abc.Codec, dict] = dict()):
+        assert(len(data) > 0)
+        is_zarr = (self.backend == 'zarr')
+
+        curr_len = self.n_steps
+        episode_length = None
+        for key, value in data.items():
+            assert(len(value.shape) >= 1)
+            if episode_length is None:
+                episode_length = len(value)
+            else:
+                assert(episode_length == len(value))
+        new_len = curr_len + episode_length
+
+        for key, value in data.items():
+            new_shape = (new_len,) + value.shape[1:]
+            # create array
+            if key not in self.data:
+                if is_zarr:
+                    cks = self._resolve_array_chunks(
+                        chunks=chunks, key=key, array=value)
+                    cpr = self._resolve_array_compressor(
+                        compressors=compressors, key=key, array=value)
+                    arr = self.data.zeros(name=key, 
+                        shape=new_shape, 
+                        chunks=cks,
+                        dtype=value.dtype,
+                        compressor=cpr)
+                else:
+                    # copy data to prevent modify
+                    arr = np.zeros(shape=new_shape, dtype=value.dtype)
+                    self.data[key] = arr
+            else:
+                arr = self.data[key]
+                assert(value.shape[1:] == arr.shape[1:])
+                # same method for both zarr and numpy
+                if is_zarr:
+                    arr.resize(new_shape)
+                else:
+                    arr.resize(new_shape, refcheck=False)
+            # copy data
+            arr[-value.shape[0]:] = value
+
+        # append to episode ends
+        episode_ends = self.episode_ends
+        if is_zarr:
+            episode_ends.resize(episode_ends.shape[0] + 1)
+        else:
+            episode_ends.resize(episode_ends.shape[0] + 1, refcheck=False)
+        episode_ends[-1] = new_len
+
+        # add meta data, e.g., votes
+        if meta_data is not None:
+            for key, value in meta_data.items():
+                if key in self.meta:
+                    meta_arr = self.meta[key]
+                    if is_zarr:
+                        meta_arr.resize(meta_arr.shape[0] + 1)
+                    else:
+                        meta_arr.resize(meta_arr.shape[0] + 1, refcheck=False)
+                    meta_arr[-1] = value
+
+        # rechunk
+        if is_zarr:
+            if episode_ends.chunks[0] < episode_ends.shape[0]:
+                rechunk_recompress_array(self.meta, 'episode_ends', 
+                    chunk_length=int(episode_ends.shape[0] * 1.5))
+    
+    def drop_episode(self):
+        is_zarr = (self.backend == 'zarr')
+        episode_ends = self.episode_ends[:].copy()
+        assert(len(episode_ends) > 0)
+        start_idx = 0
+        if len(episode_ends) > 1:
+            start_idx = episode_ends[-2]
+        for key, value in self.data.items():
+            new_shape = (start_idx,) + value.shape[1:]
+            if is_zarr:
+                value.resize(new_shape)
+            else:
+                value.resize(new_shape, refcheck=False)
+        if is_zarr:
+            self.episode_ends.resize(len(episode_ends)-1)
+        else:
+            self.episode_ends.resize(len(episode_ends)-1, refcheck=False)
+    
+    def pop_episode(self):
+        assert(self.n_episodes > 0)
+        episode = self.get_episode(self.n_episodes-1, copy=True)
+        self.drop_episode()
+        return episode
+
+    def extend(self, data):
+        self.add_episode(data)
+
+    def get_episode(self, idx, keys=None, copy=False):
+        idx = list(range(len(self.episode_ends)))[idx]
+        start_idx = 0
+        if idx > 0:
+            start_idx = self.episode_ends[idx-1]
+        end_idx = self.episode_ends[idx]
+        # result = self.get_steps_slice(start_idx, end_idx, copy=copy)
+
+        if keys is None:
+            keys = self.keys()
+
+        result = dict()
+        for key in keys:
+            value = self.data[key]
+            x = value[start_idx:end_idx]
+            if copy and isinstance(value, np.ndarray):
+                x = x.copy()
+            result[key] = x
+        return result
+    
+    def get_episode_slice(self, idx):
+        start_idx = 0
+        if idx > 0:
+            start_idx = self.episode_ends[idx-1]
+        end_idx = self.episode_ends[idx]
+        return slice(start_idx, end_idx)
+
+    def get_steps_slice(self, start, stop, step=None, copy=False):
+        _slice = slice(start, stop, step)
+
+        result = dict()
+        for key, value in self.data.items():
+            x = value[_slice]
+            if copy and isinstance(value, np.ndarray):
+                x = x.copy()
+            result[key] = x
+        return result
+    
+    # =========== chunking =============
+    def get_chunks(self) -> dict:
+        assert self.backend == 'zarr'
+        chunks = dict()
+        for key, value in self.data.items():
+            chunks[key] = value.chunks
+        return chunks
+    
+    def set_chunks(self, chunks: dict):
+        assert self.backend == 'zarr'
+        for key, value in chunks.items():
+            if key in self.data:
+                arr = self.data[key]
+                if value != arr.chunks:
+                    check_chunks_compatible(chunks=value, shape=arr.shape)
+                    rechunk_recompress_array(self.data, key, chunks=value)
+
+    def get_compressors(self) -> dict:
+        assert self.backend == 'zarr'
+        compressors = dict()
+        for key, value in self.data.items():
+            compressors[key] = value.compressor
+        return compressors
+
+    def set_compressors(self, compressors: dict):
+        assert self.backend == 'zarr'
+        for key, value in compressors.items():
+            if key in self.data:
+                arr = self.data[key]
+                compressor = self.resolve_compressor(value)
+                if compressor != arr.compressor:
+                    rechunk_recompress_array(self.data, key, compressor=compressor)
+
+    def get_votes(self) -> np.ndarray:
+        return self.meta['votes'][:]
+

common/robomimic_config_util.py

@@ -0,0 +1,47 @@
+from omegaconf import OmegaConf
+from robomimic.config import config_factory
+import robomimic.scripts.generate_paper_configs as gpc
+from robomimic.scripts.generate_paper_configs import (
+    modify_config_for_default_image_exp,
+    modify_config_for_default_low_dim_exp,
+    modify_config_for_dataset,
+)
+
+def get_robomimic_config(
+        algo_name='bc_rnn', 
+        hdf5_type='low_dim', 
+        task_name='square', 
+        dataset_type='ph'
+    ):
+    base_dataset_dir = '/tmp/null'
+    filter_key = None
+
+    # decide whether to use low-dim or image training defaults
+    modifier_for_obs = modify_config_for_default_image_exp
+    if hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"]:
+        modifier_for_obs = modify_config_for_default_low_dim_exp
+
+    algo_config_name = "bc" if algo_name == "bc_rnn" else algo_name
+    config = config_factory(algo_name=algo_config_name)
+    # turn into default config for observation modalities (e.g.: low-dim or rgb)
+    config = modifier_for_obs(config)
+    # add in config based on the dataset
+    config = modify_config_for_dataset(
+        config=config, 
+        task_name=task_name, 
+        dataset_type=dataset_type, 
+        hdf5_type=hdf5_type, 
+        base_dataset_dir=base_dataset_dir,
+        filter_key=filter_key,
+    )
+    # add in algo hypers based on dataset
+    algo_config_modifier = getattr(gpc, f'modify_{algo_name}_config_for_dataset')
+    config = algo_config_modifier(
+        config=config, 
+        task_name=task_name, 
+        dataset_type=dataset_type, 
+        hdf5_type=hdf5_type,
+    )
+    return config
+    
+

common/robomimic_util.py

@@ -0,0 +1,177 @@
+import numpy as np
+import copy
+
+import h5py
+import robomimic.utils.obs_utils as ObsUtils
+import robomimic.utils.file_utils as FileUtils
+import robomimic.utils.env_utils as EnvUtils
+from scipy.spatial.transform import Rotation
+
+from robomimic.config import config_factory
+
+
+class RobomimicAbsoluteActionConverter:
+    def __init__(self, dataset_path, algo_name='bc'):
+        # default BC config
+        config = config_factory(algo_name=algo_name)
+
+        # read config to set up metadata for observation modalities (e.g. detecting rgb observations)
+        # must ran before create dataset
+        ObsUtils.initialize_obs_utils_with_config(config)
+
+        env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path)
+        abs_env_meta = copy.deepcopy(env_meta)
+        abs_env_meta['env_kwargs']['controller_configs']['control_delta'] = False
+
+        env = EnvUtils.create_env_from_metadata(
+            env_meta=env_meta,
+            render=False, 
+            render_offscreen=False,
+            use_image_obs=False, 
+        )
+        assert len(env.env.robots) in (1, 2)
+
+        abs_env = EnvUtils.create_env_from_metadata(
+            env_meta=abs_env_meta,
+            render=False, 
+            render_offscreen=False,
+            use_image_obs=False, 
+        )
+        assert not abs_env.env.robots[0].controller.use_delta
+
+        self.env = env
+        self.abs_env = abs_env
+        self.file = h5py.File(dataset_path, 'r')
+    
+    def __len__(self):
+        return len(self.file['data'])
+
+    def convert_actions(self, 
+            states: np.ndarray, 
+            actions: np.ndarray) -> np.ndarray:
+        """
+        Given state and delta action sequence
+        generate equivalent goal position and orientation for each step
+        keep the original gripper action intact.
+        """
+        # in case of multi robot
+        # reshape (N,14) to (N,2,7)
+        # or (N,7) to (N,1,7)
+        stacked_actions = actions.reshape(*actions.shape[:-1],-1,7)
+
+        env = self.env
+        # generate abs actions
+        action_goal_pos = np.zeros(
+            stacked_actions.shape[:-1]+(3,), 
+            dtype=stacked_actions.dtype)
+        action_goal_ori = np.zeros(
+            stacked_actions.shape[:-1]+(3,), 
+            dtype=stacked_actions.dtype)
+        action_gripper = stacked_actions[...,[-1]]
+        for i in range(len(states)):
+            _ = env.reset_to({'states': states[i]})
+
+            # taken from robot_env.py L#454
+            for idx, robot in enumerate(env.env.robots):
+                # run controller goal generator
+                robot.control(stacked_actions[i,idx], policy_step=True)
+            
+                # read pos and ori from robots
+                controller = robot.controller
+                action_goal_pos[i,idx] = controller.goal_pos
+                action_goal_ori[i,idx] = Rotation.from_matrix(
+                    controller.goal_ori).as_rotvec()
+
+        stacked_abs_actions = np.concatenate([
+            action_goal_pos,
+            action_goal_ori,
+            action_gripper
+        ], axis=-1)
+        abs_actions = stacked_abs_actions.reshape(actions.shape)
+        return abs_actions
+
+    def convert_idx(self, idx):
+        file = self.file
+        demo = file[f'data/demo_{idx}']
+        # input
+        states = demo['states'][:]
+        actions = demo['actions'][:]
+
+        # generate abs actions
+        abs_actions = self.convert_actions(states, actions)
+        return abs_actions
+
+    def convert_and_eval_idx(self, idx):
+        env = self.env
+        abs_env = self.abs_env
+        file = self.file
+        # first step have high error for some reason, not representative
+        eval_skip_steps = 1
+
+        demo = file[f'data/demo_{idx}']
+        # input
+        states = demo['states'][:]
+        actions = demo['actions'][:]
+
+        # generate abs actions
+        abs_actions = self.convert_actions(states, actions)
+
+        # verify
+        robot0_eef_pos = demo['obs']['robot0_eef_pos'][:]
+        robot0_eef_quat = demo['obs']['robot0_eef_quat'][:]
+
+        delta_error_info = self.evaluate_rollout_error(
+            env, states, actions, robot0_eef_pos, robot0_eef_quat, 
+            metric_skip_steps=eval_skip_steps)
+        abs_error_info = self.evaluate_rollout_error(
+            abs_env, states, abs_actions, robot0_eef_pos, robot0_eef_quat,
+            metric_skip_steps=eval_skip_steps)
+
+        info = {
+            'delta_max_error': delta_error_info,
+            'abs_max_error': abs_error_info
+        }
+        return abs_actions, info
+
+    @staticmethod
+    def evaluate_rollout_error(env, 
+            states, actions, 
+            robot0_eef_pos, 
+            robot0_eef_quat, 
+            metric_skip_steps=1):
+        # first step have high error for some reason, not representative
+
+        # evaluate abs actions
+        rollout_next_states = list()
+        rollout_next_eef_pos = list()
+        rollout_next_eef_quat = list()
+        obs = env.reset_to({'states': states[0]})
+        for i in range(len(states)):
+            obs = env.reset_to({'states': states[i]})
+            obs, reward, done, info = env.step(actions[i])
+            obs = env.get_observation()
+            rollout_next_states.append(env.get_state()['states'])
+            rollout_next_eef_pos.append(obs['robot0_eef_pos'])
+            rollout_next_eef_quat.append(obs['robot0_eef_quat'])
+        rollout_next_states = np.array(rollout_next_states)
+        rollout_next_eef_pos = np.array(rollout_next_eef_pos)
+        rollout_next_eef_quat = np.array(rollout_next_eef_quat)
+
+        next_state_diff = states[1:] - rollout_next_states[:-1]
+        max_next_state_diff = np.max(np.abs(next_state_diff[metric_skip_steps:]))
+
+        next_eef_pos_diff = robot0_eef_pos[1:] - rollout_next_eef_pos[:-1]
+        next_eef_pos_dist = np.linalg.norm(next_eef_pos_diff, axis=-1)
+        max_next_eef_pos_dist = next_eef_pos_dist[metric_skip_steps:].max()
+
+        next_eef_rot_diff = Rotation.from_quat(robot0_eef_quat[1:]) \
+            * Rotation.from_quat(rollout_next_eef_quat[:-1]).inv()
+        next_eef_rot_dist = next_eef_rot_diff.magnitude()
+        max_next_eef_rot_dist = next_eef_rot_dist[metric_skip_steps:].max()
+
+        info = {
+            'state': max_next_state_diff,
+            'pos': max_next_eef_pos_dist,
+            'rot': max_next_eef_rot_dist
+        }
+        return info

common/sampler.py

@@ -0,0 +1,154 @@
+from typing import Optional
+import numpy as np
+import numba
+from diffusion_policy.common.replay_buffer import ReplayBuffer
+
+import pdb
+
+@numba.jit(nopython=True)
+def create_indices(
+    episode_ends:np.ndarray, sequence_length:int, 
+    episode_mask: np.ndarray,
+    pad_before: int=0, pad_after: int=0,
+    debug:bool=True) -> np.ndarray:
+    episode_mask.shape == episode_ends.shape        
+    pad_before = min(max(pad_before, 0), sequence_length-1)
+    pad_after = min(max(pad_after, 0), sequence_length-1)
+
+    indices = list()
+    for i in range(len(episode_ends)):
+        if not episode_mask[i]:
+            # skip episode
+            continue
+        start_idx = 0
+        if i > 0:
+            start_idx = episode_ends[i-1]
+        end_idx = episode_ends[i]
+        episode_length = end_idx - start_idx
+        
+        min_start = -pad_before
+        max_start = episode_length - sequence_length + pad_after
+        
+        # range stops one idx before end
+        for idx in range(min_start, max_start+1):
+            buffer_start_idx = max(idx, 0) + start_idx
+            buffer_end_idx = min(idx+sequence_length, episode_length) + start_idx
+            start_offset = buffer_start_idx - (idx+start_idx)
+            end_offset = (idx+sequence_length+start_idx) - buffer_end_idx
+            sample_start_idx = 0 + start_offset
+            sample_end_idx = sequence_length - end_offset
+            if debug:
+                assert(start_offset >= 0)
+                assert(end_offset >= 0)
+                assert (sample_end_idx - sample_start_idx) == (buffer_end_idx - buffer_start_idx)
+            indices.append([
+                buffer_start_idx, buffer_end_idx, 
+                sample_start_idx, sample_end_idx])
+    indices = np.array(indices)
+    return indices
+
+
+def get_val_mask(n_episodes, val_ratio, seed=0):
+    val_mask = np.zeros(n_episodes, dtype=bool)
+    if val_ratio <= 0:
+        return val_mask
+
+    # have at least 1 episode for validation, and at least 1 episode for train
+    n_val = min(max(1, round(n_episodes * val_ratio)), n_episodes-1)
+    rng = np.random.default_rng(seed=seed)
+    val_idxs = rng.choice(n_episodes, size=n_val, replace=False)
+    val_mask[val_idxs] = True
+    return val_mask
+
+
+def downsample_mask(mask, max_n, seed=0):
+    # subsample training data
+    train_mask = mask
+    if (max_n is not None) and (np.sum(train_mask) > max_n):
+        n_train = int(max_n)
+        curr_train_idxs = np.nonzero(train_mask)[0]
+        rng = np.random.default_rng(seed=seed)
+        train_idxs_idx = rng.choice(len(curr_train_idxs), size=n_train, replace=False)
+        train_idxs = curr_train_idxs[train_idxs_idx]
+        train_mask = np.zeros_like(train_mask)
+        train_mask[train_idxs] = True
+        assert np.sum(train_mask) == n_train
+    return train_mask
+
+class SequenceSampler:
+    def __init__(self, 
+        replay_buffer: ReplayBuffer, 
+        sequence_length:int,
+        pad_before:int=0,
+        pad_after:int=0,
+        keys=None,
+        key_first_k=dict(),
+        episode_mask: Optional[np.ndarray]=None,
+        ):
+        """
+        key_first_k: dict str: int
+            Only take first k data from these keys (to improve perf)
+        """
+
+        super().__init__()
+        assert(sequence_length >= 1)
+        if keys is None:
+            keys = list(replay_buffer.keys())
+        
+        episode_ends = replay_buffer.episode_ends[:]
+        if episode_mask is None:
+            episode_mask = np.ones(episode_ends.shape, dtype=bool)
+
+        if np.any(episode_mask):
+            indices = create_indices(episode_ends, 
+                sequence_length=sequence_length, 
+                pad_before=pad_before, 
+                pad_after=pad_after,
+                episode_mask=episode_mask
+                )
+        else:
+            indices = np.zeros((0,4), dtype=np.int64)
+
+        # (buffer_start_idx, buffer_end_idx, sample_start_idx, sample_end_idx)
+        self.indices = indices 
+        self.keys = list(keys) # prevent OmegaConf list performance problem
+        self.sequence_length = sequence_length
+        self.replay_buffer = replay_buffer
+        self.key_first_k = key_first_k
+    
+    def __len__(self):
+        return len(self.indices)
+        
+    def sample_sequence(self, idx):
+        buffer_start_idx, buffer_end_idx, sample_start_idx, sample_end_idx \
+            = self.indices[idx]
+        result = dict()
+        for key in self.keys:
+            input_arr = self.replay_buffer[key]
+            # performance optimization, avoid small allocation if possible
+            if key not in self.key_first_k:
+                sample = input_arr[buffer_start_idx:buffer_end_idx]
+            else:
+                # performance optimization, only load used obs steps
+                n_data = buffer_end_idx - buffer_start_idx
+                k_data = min(self.key_first_k[key], n_data)
+                # fill value with Nan to catch bugs
+                # the non-loaded region should never be used
+                sample = np.full((n_data,) + input_arr.shape[1:], 
+                    fill_value=np.nan, dtype=input_arr.dtype)
+                try:
+                    sample[:k_data] = input_arr[buffer_start_idx:buffer_start_idx+k_data]
+                except Exception as e:
+                    import pdb; pdb.set_trace()
+            data = sample
+            if (sample_start_idx > 0) or (sample_end_idx < self.sequence_length):
+                data = np.zeros(
+                    shape=(self.sequence_length,) + input_arr.shape[1:],
+                    dtype=input_arr.dtype)
+                if sample_start_idx > 0:
+                    data[:sample_start_idx] = sample[0]
+                if sample_end_idx < self.sequence_length:
+                    data[sample_end_idx:] = sample[-1]
+                data[sample_start_idx:sample_end_idx] = sample
+            result[key] = data
+        return result

common/timestamp_accumulator.py

@@ -0,0 +1,222 @@
+from typing import List, Tuple, Optional, Dict
+import math
+import numpy as np
+
+
+def get_accumulate_timestamp_idxs(
+    timestamps: List[float],  
+    start_time: float, 
+    dt: float, 
+    eps:float=1e-5,
+    next_global_idx: Optional[int]=0,
+    allow_negative=False
+    ) -> Tuple[List[int], List[int], int]:
+    """
+    For each dt window, choose the first timestamp in the window.
+    Assumes timestamps sorted. One timestamp might be chosen multiple times due to dropped frames.
+    next_global_idx should start at 0 normally, and then use the returned next_global_idx. 
+    However, when overwiting previous values are desired, set last_global_idx to None.
+
+    Returns:
+    local_idxs: which index in the given timestamps array to chose from
+    global_idxs: the global index of each chosen timestamp
+    next_global_idx: used for next call.
+    """
+    local_idxs = list()
+    global_idxs = list()
+    for local_idx, ts in enumerate(timestamps):
+        # add eps * dt to timestamps so that when ts == start_time + k * dt 
+        # is always recorded as kth element (avoiding floating point errors)
+        global_idx = math.floor((ts - start_time) / dt + eps)
+        if (not allow_negative) and (global_idx < 0):
+            continue
+        if next_global_idx is None:
+            next_global_idx = global_idx
+
+        n_repeats = max(0, global_idx - next_global_idx + 1)
+        for i in range(n_repeats):
+            local_idxs.append(local_idx)
+            global_idxs.append(next_global_idx + i)
+        next_global_idx += n_repeats
+    return local_idxs, global_idxs, next_global_idx
+
+
+def align_timestamps(    
+        timestamps: List[float], 
+        target_global_idxs: List[int], 
+        start_time: float, 
+        dt: float, 
+        eps:float=1e-5):
+    if isinstance(target_global_idxs, np.ndarray):
+        target_global_idxs = target_global_idxs.tolist()
+    assert len(target_global_idxs) > 0
+
+    local_idxs, global_idxs, _ = get_accumulate_timestamp_idxs(
+        timestamps=timestamps,
+        start_time=start_time,
+        dt=dt,
+        eps=eps,
+        next_global_idx=target_global_idxs[0],
+        allow_negative=True
+    )
+    if len(global_idxs) > len(target_global_idxs):
+        # if more steps available, truncate
+        global_idxs = global_idxs[:len(target_global_idxs)]
+        local_idxs = local_idxs[:len(target_global_idxs)]
+    
+    if len(global_idxs) == 0:
+        import pdb; pdb.set_trace()
+
+    for i in range(len(target_global_idxs) - len(global_idxs)):
+        # if missing, repeat
+        local_idxs.append(len(timestamps)-1)
+        global_idxs.append(global_idxs[-1] + 1)
+    assert global_idxs == target_global_idxs
+    assert len(local_idxs) == len(global_idxs)
+    return local_idxs
+
+
+class TimestampObsAccumulator:
+    def __init__(self, 
+            start_time: float, 
+            dt: float, 
+            eps: float=1e-5):
+        self.start_time = start_time
+        self.dt = dt
+        self.eps = eps
+        self.obs_buffer = dict()
+        self.timestamp_buffer = None
+        self.next_global_idx = 0
+    
+    def __len__(self):
+        return self.next_global_idx
+    
+    @property
+    def data(self):
+        if self.timestamp_buffer is None:
+            return dict()
+        result = dict()
+        for key, value in self.obs_buffer.items():
+            result[key] = value[:len(self)]
+        return result
+
+    @property
+    def actual_timestamps(self):
+        if self.timestamp_buffer is None:
+            return np.array([])
+        return self.timestamp_buffer[:len(self)]
+    
+    @property
+    def timestamps(self):
+        if self.timestamp_buffer is None:
+            return np.array([])
+        return self.start_time + np.arange(len(self)) * self.dt
+
+    def put(self, data: Dict[str, np.ndarray], timestamps: np.ndarray):
+        """
+        data:
+            key: T,*
+        """
+
+        local_idxs, global_idxs, self.next_global_idx = get_accumulate_timestamp_idxs(
+            timestamps=timestamps,
+            start_time=self.start_time,
+            dt=self.dt,
+            eps=self.eps,
+            next_global_idx=self.next_global_idx
+        )
+
+        if len(global_idxs) > 0:
+            if self.timestamp_buffer is None:
+                # first allocation
+                self.obs_buffer = dict()
+                for key, value in data.items():
+                    self.obs_buffer[key] = np.zeros_like(value)
+                self.timestamp_buffer = np.zeros(
+                    (len(timestamps),), dtype=np.float64)
+            
+            this_max_size = global_idxs[-1] + 1
+            if this_max_size > len(self.timestamp_buffer):
+                # reallocate
+                new_size = max(this_max_size, len(self.timestamp_buffer) * 2)
+                for key in list(self.obs_buffer.keys()):
+                    new_shape = (new_size,) + self.obs_buffer[key].shape[1:]
+                    self.obs_buffer[key] = np.resize(self.obs_buffer[key], new_shape)
+                self.timestamp_buffer = np.resize(self.timestamp_buffer, (new_size))
+            
+            # write data
+            for key, value in self.obs_buffer.items():
+                value[global_idxs] = data[key][local_idxs]
+            self.timestamp_buffer[global_idxs] = timestamps[local_idxs]
+
+
+class TimestampActionAccumulator:
+    def __init__(self, 
+            start_time: float, 
+            dt: float, 
+            eps: float=1e-5):
+        """
+        Different from Obs accumulator, the action accumulator
+        allows overwriting previous values.
+        """
+        self.start_time = start_time
+        self.dt = dt
+        self.eps = eps
+        self.action_buffer = None
+        self.timestamp_buffer = None
+        self.size = 0
+    
+    def __len__(self):
+        return self.size
+    
+    @property
+    def actions(self):
+        if self.action_buffer is None:
+            return np.array([])
+        return self.action_buffer[:len(self)]
+    
+    @property
+    def actual_timestamps(self):
+        if self.timestamp_buffer is None:
+            return np.array([])
+        return self.timestamp_buffer[:len(self)]
+    
+    @property
+    def timestamps(self):
+        if self.timestamp_buffer is None:
+            return np.array([])
+        return self.start_time + np.arange(len(self)) * self.dt
+
+    def put(self, actions: np.ndarray, timestamps: np.ndarray):
+        """
+        Note: timestamps is the time when the action will be issued, 
+        not when the action will be completed (target_timestamp)
+        """
+
+        local_idxs, global_idxs, _ = get_accumulate_timestamp_idxs(
+            timestamps=timestamps,
+            start_time=self.start_time,
+            dt=self.dt,
+            eps=self.eps,
+            # allows overwriting previous actions
+            next_global_idx=None
+        )
+
+        if len(global_idxs) > 0:
+            if self.timestamp_buffer is None:
+                # first allocation
+                self.action_buffer = np.zeros_like(actions)
+                self.timestamp_buffer = np.zeros((len(actions),), dtype=np.float64)
+
+            this_max_size = global_idxs[-1] + 1
+            if this_max_size > len(self.timestamp_buffer):
+                # reallocate
+                new_size = max(this_max_size, len(self.timestamp_buffer) * 2)
+                new_shape = (new_size,) + self.action_buffer.shape[1:]
+                self.action_buffer = np.resize(self.action_buffer, new_shape)
+                self.timestamp_buffer = np.resize(self.timestamp_buffer, (new_size,))
+            
+            # potentially rewrite old data (as expected)
+            self.action_buffer[global_idxs] = actions[local_idxs]
+            self.timestamp_buffer[global_idxs] = timestamps[local_idxs]
+            self.size = max(self.size, this_max_size)

config/data_collect/diffusion_policy_cnn/datacollect_diffusion_policy_cnn_lowdim_can_mh_abs.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: can_lowdim_abs
+
+name: datacollect_diffusion_unet_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 16
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_cnn/datacollect_diffusion_policy_cnn_lowdim_kitchen_abs.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: kitchen_lowdim_abs
+
+name: datacollect_diffusion_unet_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 16
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_cnn/datacollect_diffusion_policy_cnn_lowdim_lift_mh_abs.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: lift_lowdim_abs
+
+name: datacollect_diffusion_unet_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 16
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_cnn/datacollect_diffusion_policy_cnn_lowdim_pusht.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: pusht_lowdim
+
+name: datacollect_diffusion_unet_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 16
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_cnn/datacollect_diffusion_policy_cnn_lowdim_square_mh_abs.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: square_lowdim_abs
+
+name: datacollect_diffusion_unet_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 16
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_cnn/datacollect_diffusion_policy_cnn_lowdim_transport_mh_abs.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: transport_lowdim_abs
+
+name: datacollect_diffusion_unet_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 16
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_transformer/datacollect_diffusion_policy_transformer_lowdim_can_mh_abs.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: can_lowdim_abs
+
+name: datacollect_diffusion_transformer_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 10
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_transformer/datacollect_diffusion_policy_transformer_lowdim_kitchen_abs.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: kitchen_lowdim_abs
+
+name: datacollect_diffusion_transformer_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 16
+n_obs_steps: 4
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_transformer/datacollect_diffusion_policy_transformer_lowdim_lift_mh_abs.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: lift_lowdim_abs
+
+name: datacollect_diffusion_transformer_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 10
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_transformer/datacollect_diffusion_policy_transformer_lowdim_pusht.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: pusht_lowdim
+
+name: datacollect_diffusion_transformer_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 16
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_transformer/datacollect_diffusion_policy_transformer_lowdim_square_mh_abs.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: square_lowdim_abs
+
+name: datacollect_diffusion_transformer_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 10
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_transformer/datacollect_diffusion_policy_transformer_lowdim_tool_hang_ph_abs.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: tool_hang_lowdim_abs
+
+name: datacollect_diffusion_transformer_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 10
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/data_collect/diffusion_policy_transformer/datacollect_diffusion_policy_transformer_lowdim_transport_mh_abs.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+  - task: transport_lowdim_abs
+
+name: datacollect_diffusion_transformer_lowdim
+_target_: diffusion_policy.workspace.datacollect_diffusion_lowdim_workspace.DatacollectDiffusionLowdimWorkspace
+checkpoint_dir: null
+
+task_name: ${task.name}
+
+horizon: 10
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+
+collecting:
+  device: "cuda:0"
+  seed: 42
+  num_episodes: 500  # number of episodes to collect
+  render_image: True
+
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/collect_data/${task_name}/${name}_${task_name}

config/pbrl_dp_transformer/train_pbrl_diffusion_policy_transformer_lowdim_can_mh_abs.yaml

@@ -0,0 +1,163 @@
+defaults:
+  - _self_
+  - task: can_lowdim_abs_pbrl
+
+name: train_pbrl_diffusion_transformer
+
+_target_: diffusion_policy.workspace.train_pbrl_diffusion_transformer_lowdim_offline_workspace.PbrlDiffusionTransformerLowdimWorkspace
+checkpoint_dir: logs/pretrain/can_lowdim/train_diffusion_transformer/checkpoints/epoch=0040-test_mean_score=0.380.ckpt
+
+obs_dim: ${task.obs_dim}
+action_dim: ${task.action_dim}
+keypoint_dim: ${task.keypoint_dim}
+task_name: ${task.name}
+
+gamma: 0.999    # discount factor in MDP
+horizon: 10
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+obs_as_cond: True
+pred_action_steps_only: False
+
+policy:
+  _target_: diffusion_policy.policy.cpl_diffusion_transformer_lowdim_policy.CplDiffusionTransformerLowdimPolicy
+
+  model:
+    _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion
+    input_dim: ${eval:'${action_dim} if ${obs_as_cond} else ${obs_dim} + ${action_dim}'}
+    output_dim: ${policy.model.input_dim}
+    horizon: ${horizon}
+    n_obs_steps: ${n_obs_steps}
+    cond_dim: ${eval:'${obs_dim} if ${obs_as_cond} else 0'}
+
+    n_layer: 8
+    n_head: 4
+    n_emb: 256
+    p_drop_emb: 0.0
+    p_drop_attn: 0.3
+
+    causal_attn: True
+    time_as_cond: True # if false, use BERT like encoder only arch, time as input
+    obs_as_cond: ${obs_as_cond}
+    n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP
+  
+  noise_scheduler:
+    _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler
+    num_train_timesteps: 100
+    beta_start: 0.0001
+    beta_end: 0.02
+    beta_schedule: squaredcos_cap_v2
+    variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan
+    clip_sample: True # required when predict_epsilon=False
+    prediction_type: epsilon # or sample
+
+  horizon: ${horizon}
+  obs_dim: ${obs_dim}
+  action_dim: ${action_dim}
+  n_action_steps: ${n_action_steps}
+  n_obs_steps: ${n_obs_steps}
+  num_inference_steps: 100
+  obs_as_cond: ${obs_as_cond}
+  pred_action_steps_only: ${pred_action_steps_only}
+
+  beta: 0.1
+  bias_reg: 0.25
+  ignore_equal_pref: False
+  clip_margin: 1.0
+  smooth_label: 0.0
+  confidence_weight: False
+  cw_temperature: 0.03
+  unclip_win: False
+
+ema:
+  _target_: diffusion_policy.model.diffusion.ema_model.EMAModel
+  update_after_step: 0
+  inv_gamma: 1.0  # TODO: UAPA paper use : 2.0
+  power: 0.75     # TODO: UAPA paper use : 0.9
+  min_value: 0.0
+  max_value: 0.9999
+
+dataloader:
+  batch_size: 128
+  num_workers: 1
+  shuffle: True
+  pin_memory: True
+  persistent_workers: False
+
+
+optimizer:
+  learning_rate: 1.0e-4   # 1.0e-4
+  weight_decay: 1.0e-3
+  betas: [0.9, 0.95]
+
+training:
+  device: "cuda:0"
+  seed: 42
+  debug: False
+  resume: True
+  # optimization
+  lr_scheduler: cosine
+  # Transformer needs LR warmup
+  lr_warmup_steps: 0
+  num_epochs: 500
+  gradient_accumulate_every: 1
+  use_ema: True
+  # training loop control
+  # in epochs
+  rollout_every: 25
+  checkpoint_every: 25
+  val_every: 1
+  sample_every: 5
+  # steps per epoch
+  max_train_steps: null
+  max_val_steps: null
+  # misc
+  tqdm_interval_sec: 1.0
+
+  cpl_loss_type: 'cplkl'   # 'cpl'
+  use_expert_data_1: False
+  use_expert_data_2: False
+  n_epoch_sft: 0
+  sft_type: "pos"
+  stride_ratio: 0.5
+  pseudo_preference: False
+
+  dataset_1_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5.h5'
+  dataset_2_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5_2.h5'
+  preference_learning:
+    num_rounds: 1
+    equal_threshold: 0.05
+
+logging:
+  project: IL-DP-2026
+  resume: False
+  mode: online
+  name: ${name}_${training.seed}
+  tags: ["${task_name}"]
+  id: null
+  group: ${name}
+
+checkpoint:
+  topk:
+    monitor_key: test_mean_score
+    mode: max
+    k: 3
+    format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
+  save_last_ckpt: True
+  save_last_snapshot: False
+
+multi_run:
+  run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+  wandb_name_base: ${name}_${task_name}
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+  sweep:
+    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+    subdir: ${hydra.job.num}

config/pbrl_dp_transformer/train_pbrl_diffusion_policy_transformer_lowdim_lift_mh_abs.yaml

@@ -0,0 +1,163 @@
+defaults:
+  - _self_
+  - task: lift_lowdim_abs_pbrl
+
+name: train_pbrl_diffusion_transformer
+
+_target_: diffusion_policy.workspace.train_pbrl_diffusion_transformer_lowdim_offline_workspace.PbrlDiffusionTransformerLowdimWorkspace
+checkpoint_dir: logs/pretrain/lift_lowdim/train_diffusion_transformer/checkpoints/epoch=0015-test_mean_score=0.400.ckpt
+
+obs_dim: ${task.obs_dim}
+action_dim: ${task.action_dim}
+keypoint_dim: ${task.keypoint_dim}
+task_name: ${task.name}
+
+gamma: 0.999    # discount factor in MDP
+horizon: 10
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+obs_as_cond: True
+pred_action_steps_only: False
+
+policy:
+  _target_: diffusion_policy.policy.cpl_diffusion_transformer_lowdim_policy.CplDiffusionTransformerLowdimPolicy
+
+  model:
+    _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion
+    input_dim: ${eval:'${action_dim} if ${obs_as_cond} else ${obs_dim} + ${action_dim}'}
+    output_dim: ${policy.model.input_dim}
+    horizon: ${horizon}
+    n_obs_steps: ${n_obs_steps}
+    cond_dim: ${eval:'${obs_dim} if ${obs_as_cond} else 0'}
+
+    n_layer: 8
+    n_head: 4
+    n_emb: 256
+    p_drop_emb: 0.0
+    p_drop_attn: 0.3
+
+    causal_attn: True
+    time_as_cond: True # if false, use BERT like encoder only arch, time as input
+    obs_as_cond: ${obs_as_cond}
+    n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP
+  
+  noise_scheduler:
+    _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler
+    num_train_timesteps: 100
+    beta_start: 0.0001
+    beta_end: 0.02
+    beta_schedule: squaredcos_cap_v2
+    variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan
+    clip_sample: True # required when predict_epsilon=False
+    prediction_type: epsilon # or sample
+
+  horizon: ${horizon}
+  obs_dim: ${obs_dim}
+  action_dim: ${action_dim}
+  n_action_steps: ${n_action_steps}
+  n_obs_steps: ${n_obs_steps}
+  num_inference_steps: 100
+  obs_as_cond: ${obs_as_cond}
+  pred_action_steps_only: ${pred_action_steps_only}
+
+  beta: 0.1
+  bias_reg: 0.25
+  ignore_equal_pref: False
+  clip_margin: 1.0
+  smooth_label: 0.0
+  confidence_weight: False
+  cw_temperature: 0.03
+  unclip_win: False
+
+ema:
+  _target_: diffusion_policy.model.diffusion.ema_model.EMAModel
+  update_after_step: 0
+  inv_gamma: 1.0  # TODO: UAPA paper use : 2.0
+  power: 0.75     # TODO: UAPA paper use : 0.9
+  min_value: 0.0
+  max_value: 0.9999
+
+dataloader:
+  batch_size: 128
+  num_workers: 1
+  shuffle: True
+  pin_memory: True
+  persistent_workers: False
+
+
+optimizer:
+  learning_rate: 1.0e-4   # 1.0e-4
+  weight_decay: 1.0e-3
+  betas: [0.9, 0.95]
+
+training:
+  device: "cuda:0"
+  seed: 42
+  debug: False
+  resume: True
+  # optimization
+  lr_scheduler: cosine
+  # Transformer needs LR warmup
+  lr_warmup_steps: 0
+  num_epochs: 500
+  gradient_accumulate_every: 1
+  use_ema: True
+  # training loop control
+  # in epochs
+  rollout_every: 25
+  checkpoint_every: 25
+  val_every: 1
+  sample_every: 5
+  # steps per epoch
+  max_train_steps: null
+  max_val_steps: null
+  # misc
+  tqdm_interval_sec: 1.0
+
+  cpl_loss_type: 'cplkl'   # 'cpl'
+  use_expert_data_1: False
+  use_expert_data_2: False
+  n_epoch_sft: 0
+  sft_type: "pos"
+  stride_ratio: 0.5
+  pseudo_preference: False
+
+  dataset_1_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5.h5'
+  dataset_2_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5_2.h5'
+  preference_learning:
+    num_rounds: 1
+    equal_threshold: 0.05
+
+logging:
+  project: IL-DP-2026
+  resume: False
+  mode: online
+  name: ${name}_${training.seed}
+  tags: ["${task_name}"]
+  id: null
+  group: ${name}
+
+checkpoint:
+  topk:
+    monitor_key: test_mean_score
+    mode: max
+    k: 3
+    format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
+  save_last_ckpt: True
+  save_last_snapshot: False
+
+multi_run:
+  run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+  wandb_name_base: ${name}_${task_name}
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+  sweep:
+    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+    subdir: ${hydra.job.num}

config/pbrl_dp_transformer/train_pbrl_diffusion_policy_transformer_lowdim_square_mh_abs.yaml

@@ -0,0 +1,163 @@
+defaults:
+  - _self_
+  - task: square_lowdim_abs_pbrl
+
+name: train_pbrl_diffusion_transformer
+
+_target_: diffusion_policy.workspace.train_pbrl_diffusion_transformer_lowdim_offline_workspace.PbrlDiffusionTransformerLowdimWorkspace
+checkpoint_dir: logs/pretrain/square_lowdim/train_diffusion_transformer/checkpoints/epoch=0450-test_mean_score=0.520.ckpt
+
+obs_dim: ${task.obs_dim}
+action_dim: ${task.action_dim}
+keypoint_dim: ${task.keypoint_dim}
+task_name: ${task.name}
+
+gamma: 0.999    # discount factor in MDP
+horizon: 10
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+obs_as_cond: True
+pred_action_steps_only: False
+
+policy:
+  _target_: diffusion_policy.policy.cpl_diffusion_transformer_lowdim_policy.CplDiffusionTransformerLowdimPolicy
+
+  model:
+    _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion
+    input_dim: ${eval:'${action_dim} if ${obs_as_cond} else ${obs_dim} + ${action_dim}'}
+    output_dim: ${policy.model.input_dim}
+    horizon: ${horizon}
+    n_obs_steps: ${n_obs_steps}
+    cond_dim: ${eval:'${obs_dim} if ${obs_as_cond} else 0'}
+
+    n_layer: 8
+    n_head: 4
+    n_emb: 256
+    p_drop_emb: 0.0
+    p_drop_attn: 0.3
+
+    causal_attn: True
+    time_as_cond: True # if false, use BERT like encoder only arch, time as input
+    obs_as_cond: ${obs_as_cond}
+    n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP
+  
+  noise_scheduler:
+    _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler
+    num_train_timesteps: 100
+    beta_start: 0.0001
+    beta_end: 0.02
+    beta_schedule: squaredcos_cap_v2
+    variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan
+    clip_sample: True # required when predict_epsilon=False
+    prediction_type: epsilon # or sample
+
+  horizon: ${horizon}
+  obs_dim: ${obs_dim}
+  action_dim: ${action_dim}
+  n_action_steps: ${n_action_steps}
+  n_obs_steps: ${n_obs_steps}
+  num_inference_steps: 100
+  obs_as_cond: ${obs_as_cond}
+  pred_action_steps_only: ${pred_action_steps_only}
+
+  beta: 0.1
+  bias_reg: 0.25
+  ignore_equal_pref: False
+  clip_margin: 1.0
+  smooth_label: 0.0
+  confidence_weight: False
+  cw_temperature: 0.03
+  unclip_win: False
+
+ema:
+  _target_: diffusion_policy.model.diffusion.ema_model.EMAModel
+  update_after_step: 0
+  inv_gamma: 1.0  # TODO: UAPA paper use : 2.0
+  power: 0.75     # TODO: UAPA paper use : 0.9
+  min_value: 0.0
+  max_value: 0.9999
+
+dataloader:
+  batch_size: 128
+  num_workers: 1
+  shuffle: True
+  pin_memory: True
+  persistent_workers: False
+
+
+optimizer:
+  learning_rate: 1.0e-4   # 1.0e-4
+  weight_decay: 1.0e-3
+  betas: [0.9, 0.95]
+
+training:
+  device: "cuda:0"
+  seed: 42
+  debug: False
+  resume: True
+  # optimization
+  lr_scheduler: cosine
+  # Transformer needs LR warmup
+  lr_warmup_steps: 0
+  num_epochs: 500
+  gradient_accumulate_every: 1
+  use_ema: True
+  # training loop control
+  # in epochs
+  rollout_every: 25
+  checkpoint_every: 25
+  val_every: 1
+  sample_every: 5
+  # steps per epoch
+  max_train_steps: null
+  max_val_steps: null
+  # misc
+  tqdm_interval_sec: 1.0
+
+  cpl_loss_type: 'cplkl'   # 'cpl'
+  use_expert_data_1: False
+  use_expert_data_2: False
+  n_epoch_sft: 0
+  sft_type: "pos"
+  stride_ratio: 0.5
+  pseudo_preference: False
+
+  dataset_1_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5.h5'
+  dataset_2_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5_2.h5'
+  preference_learning:
+    num_rounds: 1
+    equal_threshold: 0.05
+
+logging:
+  project: IL-DP-2026
+  resume: False
+  mode: online
+  name: ${name}_${training.seed}
+  tags: ["${task_name}"]
+  id: null
+  group: ${name}
+
+checkpoint:
+  topk:
+    monitor_key: test_mean_score
+    mode: max
+    k: 3
+    format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
+  save_last_ckpt: True
+  save_last_snapshot: False
+
+multi_run:
+  run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+  wandb_name_base: ${name}_${task_name}
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+  sweep:
+    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+    subdir: ${hydra.job.num}

config/pbrl_dp_transformer/train_pbrl_diffusion_policy_transformer_lowdim_tool_hang_ph_abs.yaml

@@ -0,0 +1,163 @@
+defaults:
+  - _self_
+  - task: tool_hang_lowdim_abs_pbrl
+
+name: train_pbrl_diffusion_transformer
+
+_target_: diffusion_policy.workspace.train_pbrl_diffusion_transformer_lowdim_offline_workspace.PbrlDiffusionTransformerLowdimWorkspace
+checkpoint_dir: logs/pretrain/tool_hang_lowdim/train_diffusion_transformer/checkpoints/epoch=0500-test_mean_score=0.440.ckpt
+
+obs_dim: ${task.obs_dim}
+action_dim: ${task.action_dim}
+keypoint_dim: ${task.keypoint_dim}
+task_name: ${task.name}
+
+gamma: 0.999    # discount factor in MDP
+horizon: 10
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+obs_as_cond: True
+pred_action_steps_only: False
+
+policy:
+  _target_: diffusion_policy.policy.cpl_diffusion_transformer_lowdim_policy.CplDiffusionTransformerLowdimPolicy
+
+  model:
+    _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion
+    input_dim: ${eval:'${action_dim} if ${obs_as_cond} else ${obs_dim} + ${action_dim}'}
+    output_dim: ${policy.model.input_dim}
+    horizon: ${horizon}
+    n_obs_steps: ${n_obs_steps}
+    cond_dim: ${eval:'${obs_dim} if ${obs_as_cond} else 0'}
+
+    n_layer: 8
+    n_head: 4
+    n_emb: 256
+    p_drop_emb: 0.0
+    p_drop_attn: 0.3
+
+    causal_attn: True
+    time_as_cond: True # if false, use BERT like encoder only arch, time as input
+    obs_as_cond: ${obs_as_cond}
+    n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP
+  
+  noise_scheduler:
+    _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler
+    num_train_timesteps: 100
+    beta_start: 0.0001
+    beta_end: 0.02
+    beta_schedule: squaredcos_cap_v2
+    variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan
+    clip_sample: True # required when predict_epsilon=False
+    prediction_type: epsilon # or sample
+
+  horizon: ${horizon}
+  obs_dim: ${obs_dim}
+  action_dim: ${action_dim}
+  n_action_steps: ${n_action_steps}
+  n_obs_steps: ${n_obs_steps}
+  num_inference_steps: 100
+  obs_as_cond: ${obs_as_cond}
+  pred_action_steps_only: ${pred_action_steps_only}
+
+  beta: 0.1
+  bias_reg: 0.25
+  ignore_equal_pref: False
+  clip_margin: 1.0
+  smooth_label: 0.0
+  confidence_weight: False
+  cw_temperature: 0.03
+  unclip_win: False
+
+ema:
+  _target_: diffusion_policy.model.diffusion.ema_model.EMAModel
+  update_after_step: 0
+  inv_gamma: 1.0  # TODO: UAPA paper use : 2.0
+  power: 0.75     # TODO: UAPA paper use : 0.9
+  min_value: 0.0
+  max_value: 0.9999
+
+dataloader:
+  batch_size: 128
+  num_workers: 1
+  shuffle: True
+  pin_memory: True
+  persistent_workers: False
+
+
+optimizer:
+  learning_rate: 1.0e-4   # 1.0e-4
+  weight_decay: 1.0e-3
+  betas: [0.9, 0.95]
+
+training:
+  device: "cuda:0"
+  seed: 42
+  debug: False
+  resume: True
+  # optimization
+  lr_scheduler: cosine
+  # Transformer needs LR warmup
+  lr_warmup_steps: 0
+  num_epochs: 500
+  gradient_accumulate_every: 1
+  use_ema: True
+  # training loop control
+  # in epochs
+  rollout_every: 25
+  checkpoint_every: 25
+  val_every: 1
+  sample_every: 5
+  # steps per epoch
+  max_train_steps: null
+  max_val_steps: null
+  # misc
+  tqdm_interval_sec: 1.0
+
+  cpl_loss_type: 'cplkl'   # 'cpl'
+  use_expert_data_1: False
+  use_expert_data_2: False
+  n_epoch_sft: 0
+  sft_type: "pos"
+  stride_ratio: 0.5
+  pseudo_preference: False
+
+  dataset_1_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5.h5'
+  dataset_2_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5_2.h5'
+  preference_learning:
+    num_rounds: 1
+    equal_threshold: 0.05
+
+logging:
+  project: IL-DP-2026
+  resume: False
+  mode: online
+  name: ${name}_${training.seed}
+  tags: ["${task_name}"]
+  id: null
+  group: ${name}
+
+checkpoint:
+  topk:
+    monitor_key: test_mean_score
+    mode: max
+    k: 3
+    format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
+  save_last_ckpt: True
+  save_last_snapshot: False
+
+multi_run:
+  run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+  wandb_name_base: ${name}_${task_name}
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+  sweep:
+    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+    subdir: ${hydra.job.num}

config/pbrl_dp_transformer/train_pbrl_diffusion_policy_transformer_lowdim_transport_mh_abs.yaml

@@ -0,0 +1,163 @@
+defaults:
+  - _self_
+  - task: transport_lowdim_abs_pbrl
+
+name: train_pbrl_diffusion_transformer
+
+_target_: diffusion_policy.workspace.train_pbrl_diffusion_transformer_lowdim_offline_workspace.PbrlDiffusionTransformerLowdimWorkspace
+checkpoint_dir: logs/pretrain/transport_lowdim/train_diffusion_transformer/checkpoints/epoch=0450-test_mean_score=0.240.ckpt
+
+obs_dim: ${task.obs_dim}
+action_dim: ${task.action_dim}
+keypoint_dim: ${task.keypoint_dim}
+task_name: ${task.name}
+
+gamma: 0.999    # discount factor in MDP
+horizon: 10
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+keypoint_visible_rate: 1.0
+obs_as_cond: True
+pred_action_steps_only: False
+
+policy:
+  _target_: diffusion_policy.policy.cpl_diffusion_transformer_lowdim_policy.CplDiffusionTransformerLowdimPolicy
+
+  model:
+    _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion
+    input_dim: ${eval:'${action_dim} if ${obs_as_cond} else ${obs_dim} + ${action_dim}'}
+    output_dim: ${policy.model.input_dim}
+    horizon: ${horizon}
+    n_obs_steps: ${n_obs_steps}
+    cond_dim: ${eval:'${obs_dim} if ${obs_as_cond} else 0'}
+
+    n_layer: 8
+    n_head: 4
+    n_emb: 256
+    p_drop_emb: 0.0
+    p_drop_attn: 0.3
+
+    causal_attn: True
+    time_as_cond: True # if false, use BERT like encoder only arch, time as input
+    obs_as_cond: ${obs_as_cond}
+    n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP
+  
+  noise_scheduler:
+    _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler
+    num_train_timesteps: 100
+    beta_start: 0.0001
+    beta_end: 0.02
+    beta_schedule: squaredcos_cap_v2
+    variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan
+    clip_sample: True # required when predict_epsilon=False
+    prediction_type: epsilon # or sample
+
+  horizon: ${horizon}
+  obs_dim: ${obs_dim}
+  action_dim: ${action_dim}
+  n_action_steps: ${n_action_steps}
+  n_obs_steps: ${n_obs_steps}
+  num_inference_steps: 100
+  obs_as_cond: ${obs_as_cond}
+  pred_action_steps_only: ${pred_action_steps_only}
+
+  beta: 0.1
+  bias_reg: 0.25
+  ignore_equal_pref: False
+  clip_margin: 1.0
+  smooth_label: 0.0
+  confidence_weight: False
+  cw_temperature: 0.03
+  unclip_win: False
+
+ema:
+  _target_: diffusion_policy.model.diffusion.ema_model.EMAModel
+  update_after_step: 0
+  inv_gamma: 1.0  # TODO: UAPA paper use : 2.0
+  power: 0.75     # TODO: UAPA paper use : 0.9
+  min_value: 0.0
+  max_value: 0.9999
+
+dataloader:
+  batch_size: 128
+  num_workers: 1
+  shuffle: True
+  pin_memory: True
+  persistent_workers: False
+
+
+optimizer:
+  learning_rate: 1.0e-4   # 1.0e-4
+  weight_decay: 1.0e-3
+  betas: [0.9, 0.95]
+
+training:
+  device: "cuda:0"
+  seed: 42
+  debug: False
+  resume: True
+  # optimization
+  lr_scheduler: cosine
+  # Transformer needs LR warmup
+  lr_warmup_steps: 0
+  num_epochs: 500
+  gradient_accumulate_every: 1
+  use_ema: True
+  # training loop control
+  # in epochs
+  rollout_every: 25
+  checkpoint_every: 25
+  val_every: 1
+  sample_every: 5
+  # steps per epoch
+  max_train_steps: null
+  max_val_steps: null
+  # misc
+  tqdm_interval_sec: 1.0
+
+  cpl_loss_type: 'cplkl'   # 'cpl'
+  use_expert_data_1: False
+  use_expert_data_2: False
+  n_epoch_sft: 0
+  sft_type: "pos"
+  stride_ratio: 0.5
+  pseudo_preference: False
+
+  dataset_1_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5.h5'
+  dataset_2_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5_2.h5'
+  preference_learning:
+    num_rounds: 1
+    equal_threshold: 0.05
+
+logging:
+  project: IL-DP-2026
+  resume: False
+  mode: online
+  name: ${name}_${training.seed}
+  tags: ["${task_name}"]
+  id: null
+  group: ${name}
+
+checkpoint:
+  topk:
+    monitor_key: test_mean_score
+    mode: max
+    k: 3
+    format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
+  save_last_ckpt: True
+  save_last_snapshot: False
+
+multi_run:
+  run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+  wandb_name_base: ${name}_${task_name}
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+  sweep:
+    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
+    subdir: ${hydra.job.num}

config/pbrl_reflow_transformer/train_pbrl_reflow_policy_transformer_lowdim_can_mh_abs.yaml

@@ -0,0 +1,145 @@
+defaults:
+  - _self_
+  - task: can_lowdim_abs_pbrl
+
+name: train_pbrl_reflow_transformer
+
+_target_: diffusion_policy.workspace.train_pbrl_reflow_transformer_lowdim_offline_workspace.PbrlReFlowTransformerLowdimWorkspace
+checkpoint_dir: logs/pretrain/can_lowdim/reflow_transformer_ERate.98/reflow_transformer_ERate.98_42/checkpoints/epoch=0100-test_mean_score=0.320.ckpt
+
+obs_dim: ${task.obs_dim}
+action_dim: ${task.action_dim}
+task_name: ${task.name}
+
+gamma: 0.999    # discount factor in MDP
+horizon: 10   # this horizon is for dataloader
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+
+# Used during evaluation (both during training and evaluation)
+test_clip_intermediate_actions: True
+test_denoising_steps: 8
+test_model_type: 'original'  # can be "ema", "original"
+
+policy:
+  _target_: diffusion_policy.policy.cpl_reflow_transformer_lowdim_policy.CplReFlowTransformerLowdimPolicy
+
+  model:
+    _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion
+    input_dim: ${action_dim}
+    output_dim: ${policy.model.input_dim}
+    horizon: ${eval:'${horizon}-${n_obs_steps}+1'}
+    n_obs_steps: ${n_obs_steps}
+    cond_dim: ${obs_dim}
+
+    n_layer: 8
+    n_head: 4
+    n_emb: 256
+    p_drop_emb: 0.0
+    p_drop_attn: 0.3
+
+    causal_attn: True
+    time_as_cond: True # if false, use BERT like encoder only arch, time as input
+    obs_as_cond: True
+    n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP
+
+  horizon: ${eval:'${horizon}-${n_obs_steps}+1'}
+  obs_dim: ${obs_dim}
+  action_dim: ${action_dim}
+  n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'}
+  n_obs_steps: ${n_obs_steps}
+  sample_t_type: uniform # beta, logitnormal
+
+  beta: 0.1
+  bias_reg: 0.25
+  ignore_equal_pref: False
+  clip_margin: null  # null = disabled; recommended 10.0 for beta=0.0003125, 0.3 for beta=0.01
+  smooth_label: 0.0
+  confidence_weight: False
+  cw_temperature: 0.03
+  unclip_win: False
+
+ema:
+  decay: 0.995
+
+dataloader:
+  batch_size: 128
+  num_workers: 1
+  shuffle: True
+  pin_memory: True
+  persistent_workers: False
+
+
+optimizer:
+  lr: 1e-4   # 1.0e-4
+  weight_decay: 1e-6
+
+training:
+  device: "cuda:0"
+  seed: 42
+  debug: False
+  resume: True
+  # optimization
+  lr_scheduler: cosine
+  # Transformer needs LR warmup
+  lr_warmup_steps: 0
+  epoch_start_ema: 0
+  update_ema_freq: 10
+  num_epochs: 1000
+  use_ema: True
+  # training loop control
+  # in epochs
+  rollout_every: 25
+  checkpoint_every: 25
+  val_every: 1
+  max_train_steps: null
+  max_val_steps: null
+  # misc
+  tqdm_interval_sec: 1.0
+
+  cpl_loss_type: 'cplkl'   # 'cpl'
+  use_expert_data_1: False
+  use_expert_data_2: False
+  n_epoch_sft: 0
+  sft_type: "pos"
+  stride_ratio: 0.5
+  pseudo_preference: False
+
+  dataset_1_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5.h5'
+  dataset_2_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5_2.h5'
+  preference_learning:
+    num_rounds: 1
+    equal_threshold: 0.05
+
+logging:
+  project: IL-DP-2026
+  resume: False
+  mode: online
+  name: ${name}_${training.seed}
+  tags: ["${task_name}"]
+  id: null
+  group: ${name}
+
+checkpoint:
+  topk:
+    monitor_key: test_mean_score
+    mode: max
+    k: 3
+    format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
+  save_last_ckpt: True
+  save_last_snapshot: False
+
+multi_run:
+  run_dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+  wandb_name_base: ${name}_${training.seed}
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+  sweep:
+    dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+    subdir: ${hydra.job.num}

config/pbrl_reflow_transformer/train_pbrl_reflow_policy_transformer_lowdim_lift_mh_abs.yaml

@@ -0,0 +1,145 @@
+defaults:
+  - _self_
+  - task: lift_lowdim_abs_pbrl
+
+name: train_pbrl_reflow_transformer
+
+_target_: diffusion_policy.workspace.train_pbrl_reflow_transformer_lowdim_offline_workspace.PbrlReFlowTransformerLowdimWorkspace
+checkpoint_dir: logs/pretrain/lift_lowdim/reflow_transformer_ERate.98/reflow_transformer_ERate.98_42/checkpoints/epoch=0005-test_mean_score=0.400.ckpt
+
+obs_dim: ${task.obs_dim}
+action_dim: ${task.action_dim}
+task_name: ${task.name}
+
+gamma: 0.999    # discount factor in MDP
+horizon: 10   # this horizon is for dataloader
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+
+# Used during evaluation (both during training and evaluation)
+test_clip_intermediate_actions: True
+test_denoising_steps: 8
+test_model_type: 'original'  # can be "ema", "original"
+
+policy:
+  _target_: diffusion_policy.policy.cpl_reflow_transformer_lowdim_policy.CplReFlowTransformerLowdimPolicy
+
+  model:
+    _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion
+    input_dim: ${action_dim}
+    output_dim: ${policy.model.input_dim}
+    horizon: ${eval:'${horizon}-${n_obs_steps}+1'}
+    n_obs_steps: ${n_obs_steps}
+    cond_dim: ${obs_dim}
+
+    n_layer: 8
+    n_head: 4
+    n_emb: 256
+    p_drop_emb: 0.0
+    p_drop_attn: 0.3
+
+    causal_attn: True
+    time_as_cond: True # if false, use BERT like encoder only arch, time as input
+    obs_as_cond: True
+    n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP
+
+  horizon: ${eval:'${horizon}-${n_obs_steps}+1'}
+  obs_dim: ${obs_dim}
+  action_dim: ${action_dim}
+  n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'}
+  n_obs_steps: ${n_obs_steps}
+  sample_t_type: uniform # beta, logitnormal
+
+  beta: 0.1
+  bias_reg: 0.25
+  ignore_equal_pref: False
+  clip_margin: null  # null = disabled; recommended 10.0 for beta=0.0003125, 0.3 for beta=0.01
+  smooth_label: 0.0
+  confidence_weight: False
+  cw_temperature: 0.03
+  unclip_win: False
+
+ema:
+  decay: 0.995
+
+dataloader:
+  batch_size: 128
+  num_workers: 1
+  shuffle: True
+  pin_memory: True
+  persistent_workers: False
+
+
+optimizer:
+  lr: 1e-4   # 1.0e-4
+  weight_decay: 1e-6
+
+training:
+  device: "cuda:0"
+  seed: 42
+  debug: False
+  resume: True
+  # optimization
+  lr_scheduler: cosine
+  # Transformer needs LR warmup
+  lr_warmup_steps: 0
+  epoch_start_ema: 0
+  update_ema_freq: 10
+  num_epochs: 1000
+  use_ema: True
+  # training loop control
+  # in epochs
+  rollout_every: 25
+  checkpoint_every: 25
+  val_every: 1
+  max_train_steps: null
+  max_val_steps: null
+  # misc
+  tqdm_interval_sec: 1.0
+
+  cpl_loss_type: 'cplkl'   # 'cpl'
+  use_expert_data_1: False
+  use_expert_data_2: False
+  n_epoch_sft: 0
+  sft_type: "pos"
+  stride_ratio: 0.5
+  pseudo_preference: False
+
+  dataset_1_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5.h5'
+  dataset_2_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5_2.h5'
+  preference_learning:
+    num_rounds: 1
+    equal_threshold: 0.05
+
+logging:
+  project: IL-DP-2026
+  resume: False
+  mode: online
+  name: ${name}_${training.seed}
+  tags: ["${task_name}"]
+  id: null
+  group: ${name}
+
+checkpoint:
+  topk:
+    monitor_key: test_mean_score
+    mode: max
+    k: 3
+    format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
+  save_last_ckpt: True
+  save_last_snapshot: False
+
+multi_run:
+  run_dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+  wandb_name_base: ${name}_${training.seed}
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+  sweep:
+    dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+    subdir: ${hydra.job.num}

config/pbrl_reflow_transformer/train_pbrl_reflow_policy_transformer_lowdim_square_mh_abs.yaml

@@ -0,0 +1,145 @@
+defaults:
+  - _self_
+  - task: square_lowdim_abs_pbrl
+
+name: train_pbrl_reflow_transformer
+
+_target_: diffusion_policy.workspace.train_pbrl_reflow_transformer_lowdim_offline_workspace.PbrlReFlowTransformerLowdimWorkspace
+checkpoint_dir: logs/pretrain/square_lowdim/reflow_transformer_ERate.98/reflow_transformer_ERate.98_42/checkpoints/epoch=0400-test_mean_score=0.440.ckpt
+
+obs_dim: ${task.obs_dim}
+action_dim: ${task.action_dim}
+task_name: ${task.name}
+
+gamma: 0.999    # discount factor in MDP
+horizon: 10   # this horizon is for dataloader
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+
+# Used during evaluation (both during training and evaluation)
+test_clip_intermediate_actions: True
+test_denoising_steps: 8
+test_model_type: 'original'  # can be "ema", "original"
+
+policy:
+  _target_: diffusion_policy.policy.cpl_reflow_transformer_lowdim_policy.CplReFlowTransformerLowdimPolicy
+
+  model:
+    _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion
+    input_dim: ${action_dim}
+    output_dim: ${policy.model.input_dim}
+    horizon: ${eval:'${horizon}-${n_obs_steps}+1'}
+    n_obs_steps: ${n_obs_steps}
+    cond_dim: ${obs_dim}
+
+    n_layer: 8
+    n_head: 4
+    n_emb: 256
+    p_drop_emb: 0.0
+    p_drop_attn: 0.3
+
+    causal_attn: True
+    time_as_cond: True # if false, use BERT like encoder only arch, time as input
+    obs_as_cond: True
+    n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP
+
+  horizon: ${eval:'${horizon}-${n_obs_steps}+1'}
+  obs_dim: ${obs_dim}
+  action_dim: ${action_dim}
+  n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'}
+  n_obs_steps: ${n_obs_steps}
+  sample_t_type: uniform # beta, logitnormal
+
+  beta: 0.1
+  bias_reg: 0.25
+  ignore_equal_pref: False
+  clip_margin: null  # null = disabled; recommended 10.0 for beta=0.0003125, 0.3 for beta=0.01
+  smooth_label: 0.0
+  confidence_weight: False
+  cw_temperature: 0.03
+  unclip_win: False
+
+ema:
+  decay: 0.995
+
+dataloader:
+  batch_size: 128
+  num_workers: 1
+  shuffle: True
+  pin_memory: True
+  persistent_workers: False
+
+
+optimizer:
+  lr: 1e-4   # 1.0e-4
+  weight_decay: 1e-6
+
+training:
+  device: "cuda:0"
+  seed: 42
+  debug: False
+  resume: True
+  # optimization
+  lr_scheduler: cosine
+  # Transformer needs LR warmup
+  lr_warmup_steps: 0
+  epoch_start_ema: 0
+  update_ema_freq: 10
+  num_epochs: 1000
+  use_ema: True
+  # training loop control
+  # in epochs
+  rollout_every: 25
+  checkpoint_every: 25
+  val_every: 1
+  max_train_steps: null
+  max_val_steps: null
+  # misc
+  tqdm_interval_sec: 1.0
+
+  cpl_loss_type: 'cplkl'   # 'cpl'
+  use_expert_data_1: False
+  use_expert_data_2: False
+  n_epoch_sft: 0
+  sft_type: "pos"
+  stride_ratio: 0.5
+  pseudo_preference: False
+
+  dataset_1_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5.h5'
+  dataset_2_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5_2.h5'
+  preference_learning:
+    num_rounds: 1
+    equal_threshold: 0.05
+
+logging:
+  project: IL-DP-2026
+  resume: False
+  mode: online
+  name: ${name}_${training.seed}
+  tags: ["${task_name}"]
+  id: null
+  group: ${name}
+
+checkpoint:
+  topk:
+    monitor_key: test_mean_score
+    mode: max
+    k: 3
+    format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
+  save_last_ckpt: True
+  save_last_snapshot: False
+
+multi_run:
+  run_dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+  wandb_name_base: ${name}_${training.seed}
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+  sweep:
+    dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+    subdir: ${hydra.job.num}

config/pbrl_reflow_transformer/train_pbrl_reflow_policy_transformer_lowdim_tool_hang_ph_abs.yaml

@@ -0,0 +1,145 @@
+defaults:
+  - _self_
+  - task: tool_hang_lowdim_abs_pbrl
+
+name: train_pbrl_reflow_transformer
+
+_target_: diffusion_policy.workspace.train_pbrl_reflow_transformer_lowdim_offline_workspace.PbrlReFlowTransformerLowdimWorkspace
+checkpoint_dir: logs/pretrain/tool_hang_lowdim/reflow_transformer_ERate.98/reflow_transformer_ERate.98_42/checkpoints/epoch=0600-test_mean_score=0.360.ckpt
+
+obs_dim: ${task.obs_dim}
+action_dim: ${task.action_dim}
+task_name: ${task.name}
+
+gamma: 0.999    # discount factor in MDP
+horizon: 16   # this horizon is for dataloader
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+
+# Used during evaluation (both during training and evaluation)
+test_clip_intermediate_actions: True
+test_denoising_steps: 8
+test_model_type: 'original'  # can be "ema", "original"
+
+policy:
+  _target_: diffusion_policy.policy.cpl_reflow_transformer_lowdim_policy.CplReFlowTransformerLowdimPolicy
+
+  model:
+    _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion
+    input_dim: ${action_dim}
+    output_dim: ${policy.model.input_dim}
+    horizon: ${eval:'${horizon}-${n_obs_steps}+1'}
+    n_obs_steps: ${n_obs_steps}
+    cond_dim: ${obs_dim}
+
+    n_layer: 8
+    n_head: 4
+    n_emb: 256
+    p_drop_emb: 0.0
+    p_drop_attn: 0.3
+
+    causal_attn: True
+    time_as_cond: True # if false, use BERT like encoder only arch, time as input
+    obs_as_cond: True
+    n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP
+
+  horizon: ${eval:'${horizon}-${n_obs_steps}+1'}
+  obs_dim: ${obs_dim}
+  action_dim: ${action_dim}
+  n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'}
+  n_obs_steps: ${n_obs_steps}
+  sample_t_type: uniform # beta, logitnormal
+
+  beta: 0.1
+  bias_reg: 0.25
+  ignore_equal_pref: False
+  clip_margin: null  # null = disabled; recommended 10.0 for beta=0.0003125, 0.3 for beta=0.01
+  smooth_label: 0.0
+  confidence_weight: False
+  cw_temperature: 0.03
+  unclip_win: False
+
+ema:
+  decay: 0.995
+
+dataloader:
+  batch_size: 128
+  num_workers: 1
+  shuffle: True
+  pin_memory: True
+  persistent_workers: False
+
+
+optimizer:
+  lr: 1e-4   # 1.0e-4
+  weight_decay: 1e-6
+
+training:
+  device: "cuda:0"
+  seed: 42
+  debug: False
+  resume: True
+  # optimization
+  lr_scheduler: cosine
+  # Transformer needs LR warmup
+  lr_warmup_steps: 0
+  epoch_start_ema: 0
+  update_ema_freq: 10
+  num_epochs: 1000
+  use_ema: True
+  # training loop control
+  # in epochs
+  rollout_every: 25
+  checkpoint_every: 25
+  val_every: 1
+  max_train_steps: null
+  max_val_steps: null
+  # misc
+  tqdm_interval_sec: 1.0
+
+  cpl_loss_type: 'cplkl'   # 'cpl'
+  use_expert_data_1: False
+  use_expert_data_2: False
+  n_epoch_sft: 0
+  sft_type: "pos"
+  stride_ratio: 0.5
+  pseudo_preference: False
+
+  dataset_1_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5.h5'
+  dataset_2_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5_2.h5'
+  preference_learning:
+    num_rounds: 1
+    equal_threshold: 0.05
+
+logging:
+  project: IL-DP-2026
+  resume: False
+  mode: online
+  name: ${name}_${training.seed}
+  tags: ["${task_name}"]
+  id: null
+  group: ${name}
+
+checkpoint:
+  topk:
+    monitor_key: test_mean_score
+    mode: max
+    k: 3
+    format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
+  save_last_ckpt: True
+  save_last_snapshot: False
+
+multi_run:
+  run_dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+  wandb_name_base: ${name}_${training.seed}
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+  sweep:
+    dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+    subdir: ${hydra.job.num}

config/pbrl_reflow_transformer/train_pbrl_reflow_policy_transformer_lowdim_transport_mh_abs.yaml

@@ -0,0 +1,145 @@
+defaults:
+  - _self_
+  - task: transport_lowdim_abs_pbrl
+
+name: train_pbrl_reflow_transformer
+
+_target_: diffusion_policy.workspace.train_pbrl_reflow_transformer_lowdim_offline_workspace.PbrlReFlowTransformerLowdimWorkspace
+checkpoint_dir: logs/pretrain/transport_lowdim/reflow_transformer_ERate.98/reflow_transformer_ERate.98_42/checkpoints/epoch=0300-test_mean_score=0.400.ckpt
+
+obs_dim: ${task.obs_dim}
+action_dim: ${task.action_dim}
+task_name: ${task.name}
+
+gamma: 0.999    # discount factor in MDP
+horizon: 10   # this horizon is for dataloader
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+past_action_visible: False
+
+# Used during evaluation (both during training and evaluation)
+test_clip_intermediate_actions: True
+test_denoising_steps: 8
+test_model_type: 'original'  # can be "ema", "original"
+
+policy:
+  _target_: diffusion_policy.policy.cpl_reflow_transformer_lowdim_policy.CplReFlowTransformerLowdimPolicy
+
+  model:
+    _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion
+    input_dim: ${action_dim}
+    output_dim: ${policy.model.input_dim}
+    horizon: ${eval:'${horizon}-${n_obs_steps}+1'}
+    n_obs_steps: ${n_obs_steps}
+    cond_dim: ${obs_dim}
+
+    n_layer: 8
+    n_head: 4
+    n_emb: 256
+    p_drop_emb: 0.0
+    p_drop_attn: 0.3
+
+    causal_attn: True
+    time_as_cond: True # if false, use BERT like encoder only arch, time as input
+    obs_as_cond: True
+    n_cond_layers: 0 # >0: use transformer encoder for cond, otherwise use MLP
+
+  horizon: ${eval:'${horizon}-${n_obs_steps}+1'}
+  obs_dim: ${obs_dim}
+  action_dim: ${action_dim}
+  n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'}
+  n_obs_steps: ${n_obs_steps}
+  sample_t_type: uniform # beta, logitnormal
+
+  beta: 0.1
+  bias_reg: 0.25
+  ignore_equal_pref: False
+  clip_margin: null  # null = disabled; recommended 10.0 for beta=0.0003125, 0.3 for beta=0.01
+  smooth_label: 0.0
+  confidence_weight: False
+  cw_temperature: 0.03
+  unclip_win: False
+
+ema:
+  decay: 0.995
+
+dataloader:
+  batch_size: 128
+  num_workers: 1
+  shuffle: True
+  pin_memory: True
+  persistent_workers: False
+
+
+optimizer:
+  lr: 1e-4   # 1.0e-4
+  weight_decay: 1e-6
+
+training:
+  device: "cuda:0"
+  seed: 42
+  debug: False
+  resume: True
+  # optimization
+  lr_scheduler: cosine
+  # Transformer needs LR warmup
+  lr_warmup_steps: 0
+  epoch_start_ema: 0
+  update_ema_freq: 10
+  num_epochs: 1000
+  use_ema: True
+  # training loop control
+  # in epochs
+  rollout_every: 25
+  checkpoint_every: 25
+  val_every: 1
+  max_train_steps: null
+  max_val_steps: null
+  # misc
+  tqdm_interval_sec: 1.0
+
+  cpl_loss_type: 'cplkl'   # 'cpl'
+  use_expert_data_1: False
+  use_expert_data_2: False
+  n_epoch_sft: 0
+  sft_type: "pos"
+  stride_ratio: 0.5
+  pseudo_preference: False
+
+  dataset_1_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5.h5'
+  dataset_2_dir: 'data/robomimic/datasets/lift/noraml/robomimic_data_0.5_2.h5'
+  preference_learning:
+    num_rounds: 1
+    equal_threshold: 0.05
+
+logging:
+  project: IL-DP-2026
+  resume: False
+  mode: online
+  name: ${name}_${training.seed}
+  tags: ["${task_name}"]
+  id: null
+  group: ${name}
+
+checkpoint:
+  topk:
+    monitor_key: test_mean_score
+    mode: max
+    k: 3
+    format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
+  save_last_ckpt: True
+  save_last_snapshot: False
+
+multi_run:
+  run_dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+  wandb_name_base: ${name}_${training.seed}
+
+hydra:
+  job:
+    override_dirname: ${name}
+  run:
+    dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+  sweep:
+    dir: logs/reproduce/${task_name}/${logging.group}/${logging.name}
+    subdir: ${hydra.job.num}

config/task/aloha.yaml

@@ -0,0 +1,58 @@
+name: aloha
+
+shape_meta: &shape_meta
+  # acceptable types: rgb, low_dim
+  obs:
+    high_images:
+      shape: [3, 480, 640]
+      type: rgb
+    wrist_images:
+      shape: [3, 480, 640]
+      type: rgb
+    states:
+      shape: [13]
+  action: 
+    shape: [13]
+
+task_name: &task_name aloha
+dataset_type: &dataset_type mh
+dataset_path: &dataset_path data/aloha/placing_drawer_demo/placing_drawer_demo.hdf5
+# abs_action: &abs_action True
+
+env_runner:
+  _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
+  dataset_path: *dataset_path
+  shape_meta: *shape_meta
+  n_train: 6
+  n_train_vis: 0
+  train_start_idx: 0
+  n_test: 50
+  n_test_vis: 0
+  test_start_seed: 100000
+  # use python's eval function as resolver, single-quoted string as argument
+  max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
+  n_obs_steps: ${n_obs_steps}
+  n_action_steps: ${n_action_steps}
+  render_obs_key: 'agentview_image'
+  fps: 10
+  crf: 22
+  past_action: ${past_action_visible}
+  # abs_action: *abs_action
+  tqdm_interval_sec: 1.0
+  n_envs: 25
+# evaluation at this config requires a 16 core 64GB instance.
+
+dataset:
+  _target_: diffusion_policy.dataset.aloha_replay_image_dataset.AlohaReplayImageDataset
+  shape_meta: *shape_meta
+  dataset_path: *dataset_path
+  horizon: ${horizon}
+  pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
+  pad_after: ${eval:'${n_action_steps}-1'}
+  n_obs_steps: ${dataset_obs_steps}
+  # abs_action: *abs_action
+  # rotation_rep: 'rotation_6d'
+  # use_legacy_normalizer: False
+  use_cache: True
+  seed: 42
+  val_ratio: 0.02

config/task/blockpush_lowdim_seed.yaml

@@ -0,0 +1,34 @@
+name: blockpush_lowdim_seed
+
+obs_dim: 16
+action_dim: 2
+keypoint_dim: 2
+obs_eef_target: True
+
+env_runner:
+  _target_: diffusion_policy.env_runner.blockpush_lowdim_runner.BlockPushLowdimRunner
+  n_train: 6
+  n_train_vis: 0
+  train_start_seed: 0
+  n_test: 50
+  n_test_vis: 0
+  test_start_seed: 100000
+  max_steps: 350
+  n_obs_steps: ${n_obs_steps}
+  n_action_steps: ${n_action_steps}
+  fps: 5
+  past_action: ${past_action_visible}
+  abs_action: False
+  obs_eef_target: ${task.obs_eef_target}
+  n_envs: null
+
+dataset:
+  _target_: diffusion_policy.dataset.blockpush_lowdim_dataset.BlockPushLowdimDataset
+  zarr_path: data/block_pushing/multimodal_push_seed.zarr
+  horizon: ${horizon}
+  pad_before: ${eval:'${n_obs_steps}-1'}
+  pad_after: ${eval:'${n_action_steps}-1'}
+  obs_eef_target: ${task.obs_eef_target}
+  use_manual_normalizer: False
+  seed: 42
+  val_ratio: 0.02

config/task/blockpush_lowdim_seed_abs.yaml

@@ -0,0 +1,34 @@
+name: blockpush_lowdim_seed_abs
+
+obs_dim: 16
+action_dim: 2
+keypoint_dim: 2
+obs_eef_target: True
+
+env_runner:
+  _target_: diffusion_policy.env_runner.blockpush_lowdim_runner.BlockPushLowdimRunner
+  n_train: 6
+  n_train_vis: 0
+  train_start_seed: 0
+  n_test: 50
+  n_test_vis: 0
+  test_start_seed: 100000
+  max_steps: 350
+  n_obs_steps: ${n_obs_steps}
+  n_action_steps: ${n_action_steps}
+  fps: 5
+  past_action: ${past_action_visible}
+  abs_action: True
+  obs_eef_target: ${task.obs_eef_target}
+  n_envs: null
+
+dataset:
+  _target_: diffusion_policy.dataset.blockpush_lowdim_dataset.BlockPushLowdimDataset
+  zarr_path: data/block_pushing/multimodal_push_seed_abs.zarr
+  horizon: ${horizon}
+  pad_before: ${eval:'${n_obs_steps}-1'}
+  pad_after: ${eval:'${n_action_steps}-1'}
+  obs_eef_target: ${task.obs_eef_target}
+  use_manual_normalizer: False
+  seed: 42
+  val_ratio: 0.02

config/task/can_image.yaml

@@ -0,0 +1,64 @@
+name: can_image
+
+shape_meta: &shape_meta
+  # acceptable types: rgb, low_dim
+  obs:
+    agentview_image:
+      shape: [3, 84, 84]
+      type: rgb
+    robot0_eye_in_hand_image:
+      shape: [3, 84, 84]
+      type: rgb
+    robot0_eef_pos:
+      shape: [3]
+      # type default: low_dim
+    robot0_eef_quat:
+      shape: [4]
+    robot0_gripper_qpos:
+      shape: [2]
+  action: 
+    shape: [7]
+
+task_name: &task_name can
+dataset_type: &dataset_type mh
+dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image.hdf5
+abs_action: &abs_action False
+
+env_runner:
+  _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
+  dataset_path: *dataset_path
+  shape_meta: *shape_meta
+  # costs 1GB per env
+  n_train: 6
+  n_train_vis: 0
+  train_start_idx: 0
+  n_test: 50
+  n_test_vis: 0
+  test_start_seed: 100000
+  # use python's eval function as resolver, single-quoted string as argument
+  max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
+  n_obs_steps: ${n_obs_steps}
+  n_action_steps: ${n_action_steps}
+  render_obs_key: 'agentview_image'
+  fps: 10
+  crf: 22
+  past_action: ${past_action_visible}
+  abs_action: *abs_action
+  tqdm_interval_sec: 1.0
+  n_envs: 25
+# evaluation at this config requires a 16 core 64GB instance.
+
+dataset:
+  _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
+  shape_meta: *shape_meta
+  dataset_path: *dataset_path
+  horizon: ${horizon}
+  pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
+  pad_after: ${eval:'${n_action_steps}-1'}
+  n_obs_steps: ${dataset_obs_steps}
+  abs_action: *abs_action
+  rotation_rep: 'rotation_6d'
+  use_legacy_normalizer: False
+  use_cache: True
+  seed: 42
+  val_ratio: 0.02

config/task/can_image_abs.yaml

@@ -0,0 +1,64 @@
+name: can_image
+
+shape_meta: &shape_meta
+  # acceptable types: rgb, low_dim
+  obs:
+    agentview_image:
+      shape: [3, 84, 84]
+      type: rgb
+    robot0_eye_in_hand_image:
+      shape: [3, 84, 84]
+      type: rgb
+    robot0_eef_pos:
+      shape: [3]
+      # type default: low_dim
+    robot0_eef_quat:
+      shape: [4]
+    robot0_gripper_qpos:
+      shape: [2]
+  action: 
+    shape: [10]
+
+task_name: &task_name can
+dataset_type: &dataset_type mh
+dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5
+abs_action: &abs_action True
+
+env_runner:
+  _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
+  dataset_path: *dataset_path
+  shape_meta: *shape_meta
+  # costs 1GB per env
+  n_train: 6
+  n_train_vis: 0
+  train_start_idx: 0
+  n_test: 50
+  n_test_vis: 0
+  test_start_seed: 100000
+  # use python's eval function as resolver, single-quoted string as argument
+  max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
+  n_obs_steps: ${n_obs_steps}
+  n_action_steps: ${n_action_steps}
+  render_obs_key: 'agentview_image'
+  fps: 10
+  crf: 22
+  past_action: ${past_action_visible}
+  abs_action: *abs_action
+  tqdm_interval_sec: 1.0
+  n_envs: 25
+# evaluation at this config requires a 16 core 64GB instance.
+
+dataset:
+  _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
+  shape_meta: *shape_meta
+  dataset_path: *dataset_path
+  horizon: ${horizon}
+  pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
+  pad_after: ${eval:'${n_action_steps}-1'}
+  n_obs_steps: ${dataset_obs_steps}
+  abs_action: *abs_action
+  rotation_rep: 'rotation_6d'
+  use_legacy_normalizer: False
+  use_cache: True
+  seed: 42
+  val_ratio: 0.02