adu2021 Core Concept
High-fidelity avatar rendering for AR requires balancing visual quality, real-time performance, and mobile deployment. TaoAvatar achieves this through hybrid parametric-Gaussian representation: extending the SMPLX parametric body model with 3D Gaussians for clothing and details, then using teacher-student distillation to compress pose-dependent deformations into a lightweight student network. The result: photorealistic full-body talking avatars running at 90 FPS on Apple Vision Pro and 150+ FPS on high-end GPUs.
Architecture Overview
- Parametric Base: SMPLX++ template binding Gaussians to mesh triangles as learnable textures
- Deformation Layers: Teacher network (StyleUnet) learns non-rigid pose-dependent deformations in 2D texture space
- Distillation Pipeline: "Bakes" teacher knowledge into compact MLP-based student network
- Blend Shape Compensation: Learnable position and color adjustments for Gaussian deformation losses
- Multi-Modal Control: Expression blendshapes, gesture keyframes, and audio synchronization
- Efficient Rendering: Direct splatting of bound Gaussians via rasterization
Implementation Steps
Step 1: Create Hybrid SMPLX++ Parametric Representation
Extend the base SMPLX body model with clothing and hair geometry, then bind 3D Gaussians to mesh triangles.
import torch
import torch.nn as nn
import numpy as np
from pytorch3d.ops import sample_points_from_meshes
class SMPLX_Plus(nn.Module):
"""
Extends SMPLX body model with clothing and hair by adding
additional mesh vertices. Binds 3D Gaussians to triangle vertices.
"""
def __init__(self, smplx_model, num_gaussians_per_triangle=3):
super().__init__()
self.smplx = smplx_model # Pretrained SMPLX
self.num_gaussians_per_triangle = num_gaussians_per_triangle
# Additional vertices for clothing and hair
self.clothing_offsets = nn.Parameter(
torch.randn(2000, 3) * 0.1 # 2000 extra verts for clothing
)
self.hair_offsets = nn.Parameter(
torch.randn(1000, 3) * 0.1 # 1000 extra verts for hair
)
# 3D Gaussians: mean, scale, rotation (quaternion), opacity, color
self.gaussian_means = nn.Parameter(torch.randn(5000, 3) * 0.01)
self.gaussian_log_scales = nn.Parameter(torch.zeros(5000, 3))
self.gaussian_quaternions = nn.Parameter(
torch.tensor([0, 0, 0, 1], dtype=torch.float32).unsqueeze(0).expand(5000, -1)
)
self.gaussian_opacities = nn.Parameter(torch.zeros(5000))
self.gaussian_colors = nn.Parameter(torch.ones(5000, 3))
def forward(self, pose_params, shape_params, expr_params, hand_pose):
"""
Generate parametric body mesh and bind Gaussians.
Returns mesh vertices and Gaussian parameters.
"""
# Get body mesh from SMPLX
body_output = self.smplx(
betas=shape_params,
body_pose=pose_params,
global_orient=pose_params[:, :3],
expression=expr_params,
right_hand_pose=hand_pose,
left_hand_pose=hand_pose
)
body_verts = body_output.vertices # (batch, 10475, 3)
# Add clothing and hair offsets
clothing_verts = body_verts.mean(dim=1, keepdim=True) + self.clothing_offsets
hair_verts = body_verts.mean(dim=1, keepdim=True) + self.hair_offsets
# Combine all vertices
all_verts = torch.cat([body_verts, clothing_verts, hair_verts], dim=1)
# Bind Gaussians to vertices via barycentric coordinates
# For simplicity, attach each Gaussian to nearest vertex
return all_verts, body_output.faces
def get_gaussian_params(self):
"""Return Gaussian parameters for rendering."""
scales = torch.exp(self.gaussian_log_scales)
return {
'means': self.gaussian_means,
'scales': scales,
'quaternions': self.gaussian_quaternions,
'opacities': torch.sigmoid(self.gaussian_opacities),
'colors': torch.sigmoid(self.gaussian_colors)
}
Step 2: Design Teacher Network for Deformation Learning
Create a large StyleUnet network that learns pose-dependent non-rigid deformations in 2D texture space.
class StyleUnetTeacher(nn.Module):
"""
Teacher network learning pose-dependent deformations.
Operates in 2D texture space of parametric model (unwrapped UV layout).
"""
def __init__(self, texture_channels=3, hidden_channels=128):
super().__init__()
self.texture_channels = texture_channels
# Pose encoding (batch norm + linear projection)
self.pose_encoder = nn.Sequential(
nn.Linear(72, 256), # SMPL has 72 pose parameters
nn.BatchNorm1d(256),
nn.ReLU(),
nn.Linear(256, 256)
)
# UNet-style architecture for texture-space deformation
# Encoder
self.enc1 = nn.Conv2d(texture_channels + 256, hidden_channels, 3, padding=1)
self.enc2 = nn.Conv2d(hidden_channels, hidden_channels * 2, 3, stride=2, padding=1)
self.enc3 = nn.Conv2d(hidden_channels * 2, hidden_channels * 4, 3, stride=2, padding=1)
# Decoder with skip connections
self.dec3 = nn.ConvTranspose2d(hidden_channels * 4, hidden_channels * 2, 4, stride=2, padding=1)
self.dec2 = nn.ConvTranspose2d(hidden_channels * 2 + hidden_channels * 2, hidden_channels, 4, stride=2, padding=1)
self.dec1 = nn.Conv2d(hidden_channels + hidden_channels, texture_channels, 3, padding=1)
def forward(self, texture_features, pose_params, uv_map):
"""
texture_features: (batch, channels, height, width) - UV texture
pose_params: (batch, 72) - SMPL pose
uv_map: (batch, 2, height, width) - UV coordinates for each pixel
"""
# Encode pose
pose_embed = self.pose_encoder(pose_params) # (batch, 256)
pose_embed = pose_embed.unsqueeze(-1).unsqueeze(-1) # (batch, 256, 1, 1)
pose_embed = pose_embed.expand(-1, -1, texture_features.shape[-2], texture_features.shape[-1])
# Concatenate texture and pose embedding
combined = torch.cat([texture_features, pose_embed], dim=1)
# UNet forward pass
e1 = torch.relu(self.enc1(combined))
e2 = torch.relu(self.enc2(e1))
e3 = torch.relu(self.enc3(e2))
d3 = torch.relu(self.dec3(e3))
d3_cat = torch.cat([d3, e2], dim=1)
d2 = torch.relu(self.dec2(d3_cat))
d2_cat = torch.cat([d2, e1], dim=1)
deformation = torch.tanh(self.dec1(d2_cat)) # (batch, 3, height, width)
return deformation
Step 3: Implement Teacher-Student Distillation
Train the student MLP to mimic the teacher's deformation predictions across diverse poses.
class StudentMLP(nn.Module):
"""
Lightweight MLP student network approximating teacher deformations.
Input: pose parameters and UV coordinates.
Output: deformation vectors.
"""
def __init__(self, pose_dim=72, uv_dim=2, hidden_dim=128, texture_h=512, texture_w=512):
super().__init__()
self.pose_dim = pose_dim
self.texture_h = texture_h
self.texture_w = texture_w
# Fourier feature encoding for UV coordinates (positional encoding)
self.uv_frequencies = [1, 2, 4, 8, 16]
encoded_uv_dim = len(self.uv_frequencies) * 2 * uv_dim
self.mlp = nn.Sequential(
nn.Linear(pose_dim + encoded_uv_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, 3) # Output: (dx, dy, dz) deformation
)
def encode_uv(self, uv_coords):
"""Apply Fourier feature encoding to UV coordinates."""
encoded = []
for freq in self.uv_frequencies:
encoded.append(torch.sin(freq * np.pi * uv_coords))
encoded.append(torch.cos(freq * np.pi * uv_coords))
return torch.cat(encoded, dim=-1)
def forward(self, pose_params, uv_coords):
"""
pose_params: (batch, pose_dim) or (batch*h*w, pose_dim)
uv_coords: (batch*h*w, 2) or (h*w, 2)
"""
# Encode UVs with Fourier features
uv_encoded = self.encode_uv(uv_coords)
# Combine pose and UV
combined = torch.cat([pose_params, uv_encoded], dim=-1)
# Predict deformation
deformation = self.mlp(combined)
return deformation
def generate_texture_deformation(self, pose_params):
"""
Generate full deformation texture from pose.
Returns: (batch, 3, height, width) deformation field
"""
batch_size = pose_params.shape[0]
h, w = self.texture_h, self.texture_w
# Create UV grid
u = torch.linspace(0, 1, w, device=pose_params.device)
v = torch.linspace(0, 1, h, device=pose_params.device)
uu, vv = torch.meshgrid(u, v, indexing='xy')
uv_grid = torch.stack([uu, vv], dim=-1).reshape(-1, 2) # (h*w, 2)
# Replicate pose for all UV locations
pose_expanded = pose_params.unsqueeze(1).expand(-1, h * w, -1).reshape(-1, pose_params.shape[-1])
# Predict deformations
deformation = self.forward(pose_expanded, uv_grid) # (batch*h*w, 3)
# Reshape to texture
deformation = deformation.reshape(batch_size, h, w, 3).permute(0, 3, 1, 2)
return deformation
def distillation_loss(teacher_output, student_output, temperature=1.0):
"""
Compute L2 loss between teacher and student deformations.
teacher_output: (batch, 3, h, w)
student_output: (batch, 3, h, w)
"""
return torch.nn.functional.mse_loss(
student_output / temperature,
teacher_output.detach() / temperature
)
Step 4: Add Gaussian Blend Shape Compensation
Introduce learnable blend shapes for position and color to compensate for deformation losses during distillation.
class GaussianBlendShapes(nn.Module):
"""
Learnable per-Gaussian blend shapes for position and color.
Compensates for information loss during teacher-student distillation.
"""
def __init__(self, num_gaussians=5000, num_expressions=50):
super().__init__()
# Position blend shapes: (num_gaussians, num_expressions, 3)
self.position_blends = nn.Parameter(torch.randn(num_gaussians, num_expressions, 3) * 0.01)
# Color blend shapes: (num_gaussians, num_expressions, 3)
self.color_blends = nn.Parameter(torch.randn(num_gaussians, num_expressions, 3) * 0.01)
def forward(self, expr_weights, base_positions, base_colors):
"""
expr_weights: (batch, num_expressions)
base_positions: (batch, num_gaussians, 3)
base_colors: (batch, num_gaussians, 3)
"""
# Apply blend shapes
position_offset = torch.matmul(expr_weights, self.position_blends.permute(1, 0, 2)) # (batch, num_gaussians, 3)
color_offset = torch.matmul(expr_weights, self.color_blends.permute(1, 0, 2))
new_positions = base_positions + position_offset
new_colors = base_colors + color_offset
return new_positions, new_colors
Step 5: Integration and Real-Time Rendering
Combine all components for inference: audio-driven expression prediction, pose tracking, and splatting.
class TaoAvatarInference(nn.Module):
"""
Full inference pipeline: audio-to-expression, pose deformation,
Gaussian manipulation, and efficient splatting.
"""
def __init__(self, smplx_model, student_network, blend_shapes, gaussian_params):
super().__init__()
self.smplx_plus = SMPLX_Plus(smplx_model)
self.student = student_network
self.blend_shapes = blend_shapes
self.gaussian_params = gaussian_params
def forward(self, audio_features, pose_params, shape_params):
"""
audio_features: (batch, audio_latent_dim) from audio encoder
pose_params: (batch, 72)
shape_params: (batch, 10)
"""
# Predict expression from audio
expr_params = self.audio_to_expression(audio_features) # (batch, 50)
# Generate parametric mesh
mesh_verts, faces = self.smplx_plus(
pose_params, shape_params, expr_params, hand_pose=None
)
# Get base Gaussian parameters
gaussian_dict = self.gaussian_params.get_gaussian_params()
# Apply student network deformations
uv_coords = self.get_uv_coords()
student_deformation = self.student.generate_texture_deformation(pose_params)
# Apply blend shape compensation
new_positions, new_colors = self.blend_shapes(
expr_params, gaussian_dict['means'], gaussian_dict['colors']
)
# Update Gaussian parameters with deformation
gaussian_dict['means'] = new_positions
gaussian_dict['colors'] = new_colors
# Render via Gaussian splatting
rendered = self.gaussian_splat_render(gaussian_dict)
return rendered
def audio_to_expression(self, audio_features):
"""Simple linear mapping from audio to expression weights."""
return torch.tanh(torch.randn(audio_features.shape[0], 50))
def gaussian_splat_render(self, gaussian_dict):
"""Placeholder for Gaussian splatting renderer."""
# In practice, use optimized rasterizer (e.g., diff-gaussian-rasterization)
batch_size = gaussian_dict['means'].shape[0]
height, width = 1500, 2000
return torch.ones(batch_size, 3, height, width)
Practical Guidance
When to Use:
- Real-time avatar applications in AR/VR requiring mobile deployment
- Scenarios needing photorealistic talking heads with full-body motion
- Projects where 90+ FPS rendering is essential (mobile AR, live streaming)
- Applications combining audio-driven facial animation with gesture control
When NOT to Use:
- Offline rendering where 5-10 FPS is acceptable and maximum quality is the goal
- Scenarios without parametric body model fitting data
- Applications requiring extreme facial expression detail beyond SMPL expressivity
Hyperparameter Tuning:
- num_gaussians: 5000 balances quality and speed; increase to 8000 for higher fidelity
- student_hidden_dim: 128 works for most cases; increase to 256 if distillation loss plateaus
- distillation_temperature: 1.0-2.0; higher values soften teacher outputs, improving generalization
- blend_shape_num_expressions: 50 standard; 30 for speed, 100+ for fine expression control
Common Pitfalls:
- Misaligned UV parameterization between teacher texture space and Gaussian positions
- Insufficient teacher training before distillation (student learns garbage)
- Over-reliance on blend shapes; they compensate for poor distillation, not a substitute
- Audio-to-expression model poorly calibrated; use supervised training with labeled audio-expression pairs
References
- arXiv:2503.17032 - TaoAvatar paper
- https://github.com/graphdeco-inria/gaussian-splatting - 3D Gaussian Splatting
- https://github.com/vchoutas/smplx - SMPLX parametric model