Dynamics of coarse-grained particle systems derived via the Mori-Zwanzig projection formalism commonly take the form of a (generalized) Langevin equation with configuration-dependent friction and diffusion tensors. In this article, we introduce a class of equivariant representations of tensor-valued functions based on the Atomic Cluster Expansion (ACE) framework that allows for efficient learning of such configuration-dependent friction and diffusion tensors from data. Besides satisfying the correct equivariance properties with respect to the Euclidean group E(3), the resulting heat bath models satisfy a fluctuation-dissipation relation. Moreover, our models can be extended to include additional symmetries, such as momentum conservation, to preserve the hydrodynamic properties of the particle system. We demonstrate the capabilities of the model by constructing a model of configuration-dependent tensorial electronic friction calculated from first principles that arises during reactive molecular dynamics at metal surfaces
