def predict(self, observation):
        """
        Generates a trajectory of robot arm positions given the current state.

        Args:
            observation (torch.Tensor): Current state [robot_x, robot_y, block_x, block_y, block_angle]
                                    Shape: (batch_size, 5)

        Returns:
            torch.Tensor: Sequence of (x,y) positions for the robot arm to follow
                        Shape: (batch_size, 16, 2) where:
                        - 16 is the number of steps in the trajectory
                        - 2 is the (x,y) coordinates in pixel space (0-512)

        The function first encodes the observation, then uses it to condition a diffusion
        process that gradually denoises random trajectories into smooth, purposeful movements.
        """
        observation = observation.to(self.device)
        normalized_obs = self.normalize_data(observation, self.stats["obs"])

        # encode the observation into context values for the diffusion model
        cond = self.obs_projection(self.obs_encoder(normalized_obs))
        # keeps first & second dimension sizes unchanged, and multiplies last dimension by 16
        cond = cond.view(normalized_obs.shape[0], -1, 1).expand(-1, -1, 16)

        # initialize action with noise - random noise that will be refined into a trajectory
        action = torch.randn((observation.shape[0], 2, 16), device=self.device)

        # denoise
        # at each step `t`, the current noisy trajectory (`action`) & conditioning info (context) are
        # fed into the model to predict a denoised trajectory, then uses self.noise_scheduler.step to
        # apply this prediction & slightly reduce the noise in `action` more

        self.noise_scheduler.set_timesteps(100)
        for t in self.noise_scheduler.timesteps:
            model_output = self.model(torch.cat([action, cond], dim=1), t)
            action = self.noise_scheduler.step(model_output.sample, t, action).prev_sample

        action = action.transpose(1, 2)  # reshape to [batch, 16, 2]
        action = self.unnormalize_data(action, self.stats["action"])  # scale back to coordinates
        return action

        Generates a trajectory of robot arm positions given the current state.

        Args:
            observation (torch.Tensor): Current state [robot_x, robot_y, block_x, block_y, block_angle]
                                    Shape: (batch_size, 5)

        Returns:
            torch.Tensor: Sequence of (x,y) positions for the robot arm to follow
                        Shape: (batch_size, 16, 2) where:
                        - 16 is the number of steps in the trajectory
                        - 2 is the (x,y) coordinates in pixel space (0-512)

        The function first encodes the observation, then uses it to condition a diffusion
        process that gradually denoises random trajectories into smooth, purposeful movements.
        
examples/reinforcement_learning/diffusion_policy.py