gsplat 1.0.0

data/lib/gsplat/training/image_fitter.rb

@@ -0,0 +1,158 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Training
+    # A small self-consistency image fitting loop used by examples and E2E tests.
+    class ImageFitter
+      # Image-fit metrics, optimized colors, and target image.
+      Result = Data.define(:initial_psnr, :final_psnr, :history, :colors, :target)
+
+      # Scalar image MSE used to drive the example's reverse pass.
+      class MeanSquaredError < Autograd::Function
+        class << self
+          # Evaluates scalar mean squared error.
+          # @api private
+          def forward(context, rendered, target)
+            difference = rendered - target
+            context.save(difference, rendered.size)
+            rendered.class.cast((difference**2).sum / rendered.size)
+          end
+
+          # Returns the rendered-image VJP; targets are constants.
+          # @api private
+          def backward(context, grad_output)
+            difference, count = context.saved_values
+            [(2.0 * difference / count) * grad_output.to_f, nil]
+          end
+        end
+      end
+
+      attr_reader :dtype, :height, :learning_rate, :n_gaussians, :width
+
+      # rubocop:disable Metrics/ParameterLists
+      def initialize(width: 64, height: 64, n_gaussians: 2_000, learning_rate: 20.0,
+                     seed: 42, dtype: Numo::SFloat)
+        # rubocop:enable Metrics/ParameterLists
+        @width = width
+        @height = height
+        @n_gaussians = n_gaussians
+        @learning_rate = learning_rate
+        @dtype = dtype
+        @rng = Random.new(seed)
+        validate_options!
+        build_scene
+      end
+
+      def fit(steps: 300)
+        raise ArgumentError, "steps must be a positive integer" unless steps.is_a?(Integer) && steps.positive?
+
+        target = render(@teacher_colors)
+        colors = initial_colors
+        history = [psnr(render(colors), target)]
+        steps.times do
+          variable = Autograd::Variable.new(colors, requires_grad: true)
+          rendered = render(variable)
+          MeanSquaredError.apply(rendered, target).backward
+          colors = clipped(colors - (learning_rate * variable.grad), 0.0, 1.0)
+          history << psnr(render(colors), target)
+        end
+        Result.new(
+          initial_psnr: history.first,
+          final_psnr: history.last,
+          history: history.freeze,
+          colors: colors,
+          target: target
+        )
+      end
+
+      private
+
+      def validate_options!
+        values = { width: width, height: height, n_gaussians: n_gaussians }
+        invalid = values.find { |_name, value| !value.is_a?(Integer) || !value.positive? }
+        raise ArgumentError, "#{invalid[0]} must be a positive integer" if invalid
+        raise ArgumentError, "learning_rate must be positive" unless learning_rate.positive?
+        return if [Numo::SFloat, Numo::DFloat].include?(dtype)
+
+        raise ArgumentError, "dtype must be Numo::SFloat or Numo::DFloat"
+      end
+
+      # rubocop:disable Metrics/AbcSize
+      def build_scene
+        @means = dtype.zeros(n_gaussians, 3)
+        @quaternions = dtype.zeros(n_gaussians, 4)
+        @quaternions[true, 0] = 1
+        @scales = dtype.zeros(n_gaussians, 3)
+        @opacities = dtype.ones(n_gaussians) * 0.85
+        @viewmats = dtype.zeros(1, 4, 4)
+        @viewmats[0, true, true] = dtype.eye(4)
+        focal = [width, height].max.to_f
+        @intrinsics = dtype.cast([[[focal, 0, width / 2.0], [0, focal, height / 2.0], [0, 0, 1]]])
+        side = ::Math.sqrt(n_gaussians).ceil
+        n_gaussians.times do |index|
+          grid_x = index % side
+          grid_y = index / side
+          pixel_x = ((grid_x + 0.5) * width / side) - 0.5
+          pixel_y = ((grid_y + 0.5) * height / side) - 0.5
+          depth = 2.0 + (0.2 * (index % 3))
+          @means[index, true] = [
+            (pixel_x - (width / 2.0)) * depth / focal,
+            (pixel_y - (height / 2.0)) * depth / focal,
+            depth
+          ]
+          world_scale = 0.8 * depth / focal
+          @scales[index, true] = [world_scale, world_scale, world_scale]
+        end
+        @teacher_colors = teacher_colors(side)
+      end
+      # rubocop:enable Metrics/AbcSize
+
+      def teacher_colors(side)
+        output = dtype.zeros(n_gaussians, 3)
+        n_gaussians.times do |index|
+          x_coord = (index % side).to_f / [side - 1, 1].max
+          y_coord = (index / side).to_f / [side - 1, 1].max
+          output[index, true] = [x_coord, y_coord, 0.25 + (0.5 * x_coord * y_coord)]
+        end
+        output
+      end
+
+      def initial_colors
+        values = Array.new(n_gaussians * 3) { 0.35 + (@rng.rand * 0.3) }
+        dtype.cast(values).reshape(n_gaussians, 3)
+      end
+
+      def render(colors)
+        rendered, = Gsplat.rasterization(
+          means: @means,
+          quats: @quaternions,
+          scales: @scales,
+          opacities: @opacities,
+          colors: colors,
+          viewmats: @viewmats,
+          ks: @intrinsics,
+          width: width,
+          height: height,
+          tile_size: 16
+        )
+        rendered
+      end
+
+      def psnr(rendered, target)
+        mean_squared_error = (((Ops::TensorOps.data(rendered) - target)**2).sum / target.size).to_f
+        return Float::INFINITY if mean_squared_error.zero?
+
+        -10.0 * ::Math.log10(mean_squared_error)
+      end
+
+      def clipped(values, minimum, maximum)
+        output = values.dup
+        below = output.lt(minimum)
+        above = output.gt(maximum)
+        output[below] = minimum if below.any?
+        output[above] = maximum if above.any?
+        output
+      end
+    end
+  end
+end

data/lib/gsplat/training/losses.rb

@@ -0,0 +1,191 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Training
+    # Differentiable image losses and scalar image-quality metrics.
+    module Losses
+      # Mean absolute error with a zero subgradient at exact equality.
+      class L1 < Autograd::Function
+        class << self
+          # Evaluates mean absolute error.
+          # @api private
+          def forward(context, prediction, target)
+            unless prediction.shape == target.shape
+              raise ShapeError, "L1 shapes differ: #{prediction.shape.inspect} and #{target.shape.inspect}"
+            end
+
+            difference = prediction - target
+            context.save(difference, prediction.size)
+            prediction.class.cast(difference.abs.mean)
+          end
+
+          # Propagates the L1 subgradient to prediction and target.
+          # @api private
+          def backward(context, grad_output)
+            difference, count = context.saved_values
+            signs = difference.class.zeros(*difference.shape)
+            signs[difference.gt(0)] = 1
+            signs[difference.lt(0)] = -1
+            gradient = signs * (grad_output.to_f / count)
+            [gradient, -gradient]
+          end
+        end
+      end
+
+      # Structural similarity with the standard 11x11, sigma 1.5 window.
+      class StructuralSimilarity < Autograd::Function
+        class << self
+          # Evaluates mean SSIM and stores convolution intermediates.
+          # @api private
+          def forward(context, image_a, image_b, layout:)
+            score, cache = Math::Ssim.forward(image_a, image_b, layout: layout)
+            context.save(cache)
+            image_a.class.cast(score)
+          end
+
+          # Propagates the SSIM score gradient.
+          # @api private
+          def backward(context, grad_output)
+            Math::Ssim.backward(context.saved_values.first, grad_output)
+          end
+        end
+      end
+
+      # Weighted L1 plus (1 - SSIM) reconstruction objective.
+      class Reconstruction < Autograd::Function
+        class << self
+          # Evaluates weighted L1 and SSIM reconstruction loss.
+          # @api private
+          def forward(context, prediction, target, ssim_lambda:, layout:)
+            unless prediction.shape == target.shape
+              raise ShapeError, "reconstruction shapes differ: #{prediction.shape.inspect} and #{target.shape.inspect}"
+            end
+
+            difference = prediction - target
+            score, ssim_cache = Math::Ssim.forward(prediction, target, layout: layout)
+            context.save(difference, prediction.size, ssim_cache, ssim_lambda)
+            prediction.class.cast(((1 - ssim_lambda) * difference.abs.mean.to_f) + (ssim_lambda * (1 - score)))
+          end
+
+          # Propagates the combined reconstruction gradient.
+          # @api private
+          def backward(context, grad_output)
+            difference, count, ssim_cache, ssim_lambda = context.saved_values
+            signs = difference.class.zeros(*difference.shape)
+            signs[difference.gt(0)] = 1
+            signs[difference.lt(0)] = -1
+            l1_gradient = signs * ((1 - ssim_lambda) * grad_output.to_f / count)
+            ssim_a, ssim_b = Math::Ssim.backward(ssim_cache, grad_output)
+            [l1_gradient - (ssim_lambda * ssim_a), -l1_gradient - (ssim_lambda * ssim_b)]
+          end
+        end
+      end
+
+      # Reconstruction loss with activated opacity and scale mean penalties.
+      class RegularizedReconstruction < Autograd::Function
+        class << self
+          # rubocop:disable Metrics/ParameterLists
+          def forward(context, prediction, target, opacities, scales, ssim_lambda:, opacity_reg:, scale_reg:, layout:)
+            # rubocop:enable Metrics/ParameterLists
+            difference = prediction - target
+            score, ssim_cache = Math::Ssim.forward(prediction, target, layout: layout)
+            context.save(
+              difference, prediction.size, ssim_cache, ssim_lambda,
+              opacities.shape, scales.shape, opacity_reg, scale_reg
+            )
+            reconstruction = ((1 - ssim_lambda) * difference.abs.mean.to_f) + (ssim_lambda * (1 - score))
+            prediction.class.cast(
+              reconstruction + (opacity_reg * opacities.mean.to_f) + (scale_reg * scales.mean.to_f)
+            )
+          end
+
+          # rubocop:disable Metrics/AbcSize
+          def backward(context, grad_output)
+            difference, count, cache, ssim_lambda,
+              opacity_shape, scale_shape, opacity_reg, scale_reg = context.saved_values
+            signs = difference.class.zeros(*difference.shape)
+            signs[difference.gt(0)] = 1
+            signs[difference.lt(0)] = -1
+            l1_gradient = signs * ((1 - ssim_lambda) * grad_output.to_f / count)
+            ssim_a, ssim_b = Math::Ssim.backward(cache, grad_output)
+            opacity_gradient = difference.class.ones(*opacity_shape) *
+                               opacity_reg * grad_output.to_f / opacity_shape.inject(:*)
+            scale_gradient = difference.class.ones(*scale_shape) *
+                             scale_reg * grad_output.to_f / scale_shape.inject(:*)
+            [
+              l1_gradient - (ssim_lambda * ssim_a),
+              -l1_gradient - (ssim_lambda * ssim_b),
+              opacity_gradient,
+              scale_gradient
+            ]
+          end
+          # rubocop:enable Metrics/AbcSize
+        end
+      end
+
+      module_function
+
+      # Returns mean absolute error, retaining an autograd graph when needed.
+      def l1(prediction, target)
+        return L1.apply(prediction, target) if [prediction, target].any?(Autograd::Variable)
+
+        L1.forward(Autograd::Context.new([false, false], [prediction, target]), prediction, target)
+      end
+
+      # Returns mean SSIM over batches, channels, and pixels.
+      def ssim(image_a, image_b, layout: :auto)
+        if [image_a, image_b].any?(Autograd::Variable)
+          return StructuralSimilarity.apply(image_a, image_b, layout: layout)
+        end
+
+        Math::Ssim.forward(image_a, image_b, layout: layout).first
+      end
+
+      # Returns `(1-lambda)*L1 + lambda*(1-SSIM)`.
+      def reconstruction(prediction, target, ssim_lambda: 0.2, layout: :auto)
+        raise ArgumentError, "ssim_lambda must be between 0 and 1" unless ssim_lambda.between?(0.0, 1.0)
+
+        if [prediction, target].any?(Autograd::Variable)
+          return Reconstruction.apply(
+            prediction,
+            target,
+            ssim_lambda: ssim_lambda,
+            layout: layout
+          )
+        end
+
+        l1_value = (prediction - target).abs.mean.to_f
+        ((1 - ssim_lambda) * l1_value) + (ssim_lambda * (1 - ssim(prediction, target, layout: layout)))
+      end
+
+      # Reconstruction objective plus optional MCMC opacity and scale regularizers.
+      def regularized_reconstruction(prediction, target, opacities, scales, **options)
+        RegularizedReconstruction.apply(
+          prediction,
+          target,
+          opacities,
+          scales,
+          ssim_lambda: options.fetch(:ssim_lambda, 0.2),
+          opacity_reg: options.fetch(:opacity_reg, 0.0),
+          scale_reg: options.fetch(:scale_reg, 0.0),
+          layout: options.fetch(:layout, :auto)
+        )
+      end
+
+      # Returns peak signal-to-noise ratio in decibels.
+      def psnr(prediction, target, max_value: 1.0)
+        prediction = Ops::TensorOps.data(prediction)
+        target = Ops::TensorOps.data(target)
+        unless prediction.shape == target.shape
+          raise ShapeError, "PSNR shapes differ: #{prediction.shape.inspect} and #{target.shape.inspect}"
+        end
+        raise ArgumentError, "max_value must be positive" unless max_value.positive?
+
+        mean_squared_error = ((prediction - target)**2).mean.to_f
+        return Float::INFINITY if mean_squared_error.zero?
+
+        10 * ::Math.log10((max_value**2) / mean_squared_error)
+      end
+    end
+  end
+end

data/lib/gsplat/training/scene.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+module Gsplat
+  # Dataset containers, losses, fitting utilities, and training loops.
+  module Training
+    # In-memory multi-view training scene.
+    class Scene
+      attr_reader :colors, :height, :images, :intrinsics, :names, :points,
+                  :scene_scale, :viewmats, :width
+
+      # rubocop:disable Metrics/ParameterLists
+      def initialize(viewmats:, intrinsics:, images:, points:, colors:, names: nil, scene_scale: nil)
+        # rubocop:enable Metrics/ParameterLists
+        @viewmats = viewmats
+        @intrinsics = intrinsics
+        @images = images
+        @points = points
+        @colors = colors
+        @names = names || Array.new(images.shape[0]) { |index| format("camera_%03d", index) }
+        @height = images.shape[1]
+        @width = images.shape[2]
+        validate!
+        @scene_scale = scene_scale || infer_scene_scale
+      end
+
+      # Loads a COLMAP sparse model and its corresponding RGB images.
+      #
+      # @param path [String] dataset root containing `sparse` and `images`
+      # @param data_factor [Numeric] image/intrinsics downsampling factor
+      # @return [Scene]
+      # rubocop:disable Metrics/AbcSize
+      def self.from_colmap(path, data_factor: 1)
+        dataset = IO::Colmap.read(path, data_factor: data_factor)
+        records = dataset.images.values.sort_by(&:name)
+        views = Numo::SFloat.cast(stack(records.map(&:world_to_camera)))
+        intrinsics = Numo::SFloat.cast(
+          stack(records.map { |record| dataset.cameras.fetch(record.camera_id).intrinsics })
+        )
+        image_directory = resolve_image_directory(path, data_factor)
+        images = Numo::SFloat.cast(
+          stack(records.map { |record| IO::Image.read(File.join(image_directory, record.name)) })
+        )
+        points = dataset.points3d.values.sort_by(&:id)
+        new(
+          viewmats: views,
+          intrinsics: intrinsics,
+          images: images,
+          points: Numo::SFloat.cast(stack_rows(points.map(&:xyz))),
+          colors: Numo::SFloat.cast(stack_rows(points.map { |point| Numo::DFloat.cast(point.rgb) / 255.0 })),
+          names: records.map(&:name)
+        )
+      end
+      # rubocop:enable Metrics/AbcSize
+
+      # Number of registered training views.
+      #
+      # @return [Integer]
+      def camera_count
+        images.shape[0]
+      end
+
+      class << self
+        private
+
+        def stack(arrays)
+          raise ArgumentError, "cannot stack an empty array list" if arrays.empty?
+
+          type = arrays.first.class
+          output = type.zeros(*([arrays.length] + arrays.first.shape))
+          arrays.each_with_index do |array, index|
+            output[*([index] + Array.new(array.ndim, true))] = array
+          end
+          output
+        end
+
+        def stack_rows(rows)
+          raise ArgumentError, "COLMAP point cloud is empty" if rows.empty?
+
+          rows.first.class.cast(rows.map(&:to_a))
+        end
+
+        def resolve_image_directory(path, factor)
+          scaled = File.join(path, "images_#{factor}")
+          return scaled if factor != 1 && Dir.exist?(scaled)
+
+          File.join(path, "images")
+        end
+      end
+
+      private
+
+      def validate!
+        cameras = images.shape[0]
+        valid = images.ndim == 4 && images.shape[-1] == 3 &&
+                viewmats.shape == [cameras, 4, 4] &&
+                intrinsics.shape == [cameras, 3, 3] &&
+                points.ndim == 2 && points.shape[1] == 3 &&
+                colors.shape == [points.shape[0], 3] &&
+                names.length == cameras
+        raise ShapeError, "inconsistent multi-view scene arrays" unless valid
+      end
+
+      def infer_scene_scale
+        centers = viewmats.class.zeros(viewmats.shape[0], 3)
+        viewmats.shape[0].times do |index|
+          rotation = viewmats[index, 0...3, 0...3]
+          translation = viewmats[index, 0...3, 3]
+          centers[index, true] = -translation.dot(rotation)
+        end
+        center = centers.mean(axis: 0)
+        scale = Numo::NMath.sqrt(((centers - center)**2).sum(axis: 1)).max.to_f
+        scale.positive? ? scale : 1.0
+      end
+    end
+  end
+end

data/lib/gsplat/training/trainer.rb

@@ -0,0 +1,236 @@
+# frozen_string_literal: true
+
+require "fileutils"
+
+module Gsplat
+  module Training
+    # End-to-end multi-view Gaussian training loop.
+    class Trainer
+      # Final step, before/after metrics, and periodic history.
+      Result = Data.define(:step, :initial_metrics, :final_metrics, :history)
+      # Mapping from parameter names to {Config} learning-rate attributes.
+      LEARNING_RATES = {
+        means: :means_lr,
+        scales: :scales_lr,
+        quats: :quats_lr,
+        opacities: :opacities_lr,
+        sh0: :sh0_lr,
+        shN: :shN_lr
+      }.freeze
+
+      attr_reader :config, :optimizers, :params, :scene, :step, :strategy, :strategy_state
+
+      def initialize(scene:, config: Config.new, params: nil, strategy: nil)
+        @scene = scene
+        @config = config
+        @rng = Random.new(config.seed)
+        @params = params || Utils.init_from_points(
+          scene.points,
+          scene.colors,
+          sh_degree: config.sh_degree,
+          init_opacity: config.init_opacity,
+          init_scale: config.init_scale,
+          rng: @rng
+        )
+        @optimizers = build_optimizers
+        @strategy = strategy || Strategy::Default.new
+        @strategy.check_sanity(@params, @optimizers)
+        @strategy_state = @strategy.initialize_state(scene_scale: scene.scene_scale)
+        @scheduler = Optim::ExponentialLR.new(
+          @optimizers.fetch(:means),
+          lr_final: config.means_lr_final * scene.scene_scale,
+          max_steps: config.max_steps
+        )
+        @step = 0
+      end
+
+      # Optimizes until the configured maximum or for an additional step count.
+      #
+      # @param steps [Integer, nil]
+      # @return [Result]
+      def train(steps: nil)
+        target_step = steps ? step + steps : config.max_steps
+        initial_metrics = evaluate
+        history = []
+        while step < target_step
+          @step += 1
+          loss = training_step
+          history << { step: step, loss: loss } if log_due?
+          evaluate_and_log!(history) if config.eval_steps.include?(step)
+          save_checkpoint! if config.save_steps.include?(step)
+        end
+        Result.new(
+          step: step,
+          initial_metrics: initial_metrics,
+          final_metrics: evaluate,
+          history: history.freeze
+        )
+      end
+
+      # Renders every view without graph recording and reports PSNR/SSIM.
+      #
+      # @return [Hash{Symbol=>Float}]
+      def evaluate
+        Autograd.no_grad do
+          rendered, = render((0...scene.camera_count).to_a, sh_degree_for(step))
+          values = Ops::TensorOps.data(rendered)
+          {
+            psnr: Losses.psnr(values, scene.images),
+            ssim: Losses.ssim(values, scene.images, layout: :nhwc)
+          }
+        end
+      end
+
+      # Saves current parameters, optimizer moments, step, and config.
+      #
+      # @param path [String, nil] generated from `config.output_dir` when nil
+      # @return [String] written path
+      def save_checkpoint!(path = nil)
+        path ||= File.join(config.output_dir, "checkpoints", format("step_%06d.npz", step))
+        FileUtils.mkdir_p(File.dirname(path))
+        IO::Checkpoint.save(
+          path,
+          params: params,
+          optimizers: optimizers,
+          step: step,
+          config: config.to_h
+        )
+        path
+      end
+
+      # Restores parameters, optimizer moments, step, and scheduler position.
+      #
+      # @param source [String, #read]
+      # @return [IO::Checkpoint::Snapshot]
+      def load_checkpoint!(source)
+        snapshot = IO::Checkpoint.restore!(source, params: params, optimizers: optimizers)
+        @step = snapshot.step
+        @scheduler.step(step)
+        snapshot
+      end
+
+      private
+
+      # rubocop:disable Metrics/AbcSize
+      def training_step
+        indices = batch_indices
+        rendered, _alphas, info = render(indices, sh_degree_for(step))
+        target = scene.images[indices, true, true, true].dup
+        strategy.step_pre_backward(
+          params: params, optimizers: optimizers, state: strategy_state, step: step, info: info
+        )
+        activated = activated_params
+        loss = Losses.regularized_reconstruction(
+          rendered,
+          target,
+          activated.fetch(:opacities),
+          activated.fetch(:scales),
+          ssim_lambda: config.ssim_lambda,
+          opacity_reg: config.opacity_reg,
+          scale_reg: config.scale_reg,
+          layout: :nhwc
+        )
+        loss.backward
+        strategy.step_post_backward(
+          params: params,
+          optimizers: optimizers,
+          state: strategy_state,
+          step: step,
+          info: info,
+          lr: optimizers.fetch(:means).learning_rate
+        )
+        optimizers.each_value(&:step)
+        optimizers.each_value(&:zero_grad!)
+        @scheduler.step(step)
+        loss.data.to_f
+      end
+      # rubocop:enable Metrics/AbcSize
+
+      # rubocop:disable Metrics/AbcSize
+      def render(indices, degree)
+        activated = activated_params
+        backgrounds = random_background(indices.length, scene.images.class)
+        renderer = %i[two_d 2dgs].include?(config.model_type.to_sym) ? :rasterization_2dgs : :rasterization
+        outputs = Gsplat.public_send(
+          renderer,
+          means: params.fetch(:means),
+          quats: activated.fetch(:quats),
+          scales: activated.fetch(:scales),
+          opacities: activated.fetch(:opacities),
+          colors: activated.fetch(:colors),
+          viewmats: scene.viewmats[indices, true, true].dup,
+          ks: scene.intrinsics[indices, true, true].dup,
+          width: scene.width,
+          height: scene.height,
+          sh_degree: degree,
+          backgrounds: backgrounds,
+          near_plane: config.near_plane,
+          far_plane: config.far_plane,
+          tile_size: config.tile_size,
+          rasterize_mode: config.rasterize_mode,
+          absgrad: strategy.respond_to?(:absgrad) && strategy.absgrad
+        )
+        return outputs if renderer == :rasterization
+
+        [outputs[0], outputs[1], outputs[6]]
+      end
+      # rubocop:enable Metrics/AbcSize
+
+      def activated_params
+        sh0 = params.fetch(:sh0)
+        shn = params.fetch(:shN)
+        colors = if shn.data.shape[1].zero?
+                   sh0
+                 else
+                   Ops::TensorOps.apply(Ops::ConcatCoefficients, sh0, shn)
+                 end
+        {
+          quats: Ops::TensorOps.apply(Ops::NormalizeQuaternion, params.fetch(:quats)),
+          scales: Ops::TensorOps.apply(Ops::Exp, params.fetch(:scales)),
+          opacities: Ops::TensorOps.apply(Ops::Sigmoid, params.fetch(:opacities)),
+          colors: colors
+        }
+      end
+
+      def build_optimizers
+        LEARNING_RATES.to_h do |name, config_name|
+          rate = config.public_send(config_name)
+          rate *= scene.scene_scale if name == :means
+          [name, Optim::Adam.new(params.fetch(name), lr: rate)]
+        end
+      end
+
+      def batch_indices
+        Array.new(config.batch_size) { @rng.rand(scene.camera_count) }
+      end
+
+      def sh_degree_for(current_step)
+        [config.sh_degree, current_step / config.sh_degree_interval].min
+      end
+
+      def random_background(count, type)
+        return nil unless config.random_background
+
+        type.cast(Array.new(count * 3) { @rng.rand }).reshape(count, 3)
+      end
+
+      def log_due?
+        (step % config.log_every).zero?
+      end
+
+      def evaluate_and_log!(history)
+        metrics = evaluate
+        history << { step: step, **metrics }
+        Gsplat.logger.info(
+          format(
+            "step=%<step>d psnr=%<psnr>.3f ssim=%<ssim>.5f gaussians=%<gaussians>d",
+            step: step,
+            psnr: metrics[:psnr],
+            ssim: metrics[:ssim],
+            gaussians: params.fetch(:means).data.shape[0]
+          )
+        )
+      end
+    end
+  end
+end